package main import ( "context" "flag" "fmt" "log" "net/http" "net/url" "runtime" "strings" "time" "regexp" "strconv" "encoding/json" "git.satori.love/gameui/webp/api" "git.satori.love/gameui/webp/models" _ "github.com/go-sql-driver/mysql" "github.com/graphql-go/graphql" "github.com/graphql-go/handler" "github.com/milvus-io/milvus-sdk-go/v2/entity" "github.com/spf13/viper" lru "github.com/hashicorp/golang-lru/v2" ) // string 转换为 int, 如果转换失败则返回默认值 func stringToInt(str string, defaultValue int) int { if str == "" { return defaultValue } value, err := strconv.Atoi(str) if err != nil { return defaultValue } return value } func LogComponent(startTime int64, r *http.Request) { ms := (time.Now().UnixNano() - startTime) / 1000000 color := "\033[1;32m%d\033[0m" if ms > 800 { color = "\033[1;31m%dms\033[0m" // 紅色加重 } else if ms > 500 { color = "\033[1;33m%dms\033[0m" // 黃色加重 } else if ms > 300 { color = "\033[1;32m%dms\033[0m" // 綠色加重 } else if ms > 200 { color = "\033[1;34m%dms\033[0m" // 藍色加重 } else if ms > 100 { color = "\033[1;35m%dms\033[0m" // 紫色加重 } else { color = "\033[1;36m%dms\033[0m" // 黑色加重 } endTime := fmt.Sprintf(color, ms) method := fmt.Sprintf("\033[1;32m%s\033[0m", r.Method) // 綠色加重 url := fmt.Sprintf("\033[1;34m%s\033[0m", r.URL) // 藍色加重 log.Println(method, url, endTime, r.Header.Get("X-Forwarded-For")) } func LogRequest(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 var user_id int if token := r.Header.Get("token"); token != "" { user_id = api.ParseToken(token) } next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), "user_id", user_id))) }) } type Image struct { Id int `json:"id" db:"id"` Width int `json:"width" db:"width"` Height int `json:"height" db:"height"` Content string `json:"content" db:"content"` ArticleCategoryTopId int `json:"article_category_top_id" db:"article_category_top_id"` PraiseCount int `json:"praise_count" db:"praise_count"` CollectCount int `json:"collect_count" db:"collect_count"` CreateTime time.Time `json:"createTime" db:"createTime"` UpdateTime time.Time `json:"updateTime" db:"updateTime"` UserID int `json:"user_id" db:"user_id"` User models.User `json:"user" db:"user"` Article models.Article `json:"article" db:"article"` } type Tag struct { Id int `json:"id"` Name string `json:"name"` CreateTime time.Time `json:"create_time"` UpdateTime time.Time `json:"update_time"` } type History struct { Type string `json:"type"` CreateTime time.Time `json:"create_time"` Data interface{} `json:"data"` } type ListView struct { Code int `json:"code"` Page int `json:"page"` PageSize int `json:"pageSize"` Total int `json:"total"` Next bool `json:"next"` List []interface{} `json:"list"` } var mysqlConnection models.MysqlConnection var milvusConnection models.MilvusConnection func GetNetWorkEmbedding(id int) (embedding []float32) { host := viper.GetString("embedding.host") port := viper.GetInt("embedding.port") httpClient := &http.Client{} request, err := http.NewRequest("PUT", fmt.Sprintf("http://%s:%d/api/default/%d", host, port, id), nil) if err != nil { log.Println("请求失败1:", err) return } response, err := httpClient.Do(request) if err != nil { log.Println("请求失败2:", err) return } defer response.Body.Close() var result struct { Code int `json:"code"` Message string `json:"message"` Feature []float32 `json:"feature"` } err = json.NewDecoder(response.Body).Decode(&result) if err != nil { log.Println("解析失败:", err) return } if result.Code != 0 { log.Println("请求失败3:", result.Message) return } return result.Feature } var lruCache, _ = lru.New[int, []int64](100000) func (image *Image) GetSimilarImagesIdList(collection_name string) (ids []int64) { ctx := context.Background() // 先从 LRU 中查询缓存的结果, 如果缓存中有, 直接返回 if value, ok := lruCache.Get(image.Id); ok { return value } // 先从milvus中查询图片的向量 var embedding []float32 result, err := milvusConnection.Client.Query(ctx, collection_name, nil, fmt.Sprintf("id in [%d]", image.Id), []string{"embedding"}) if err != nil { log.Println("查詢向量失敗:", err) embedding = GetNetWorkEmbedding(image.Id) } else { for _, item := range result { if item.Name() == "embedding" { embedding = item.FieldData().GetVectors().GetFloatVector().Data continue } } } // 处理向量不存在的情况 if len(embedding) == 0 { log.Println("向量不存在, 也未能重新生成") return ids } // 用向量查询相似图片 topk := 200 sp, _ := entity.NewIndexIvfFlatSearchParam(64) vectors := []entity.Vector{entity.FloatVector(embedding)} resultx, err := milvusConnection.Client.Search(ctx, collection_name, nil, "", []string{"id", "article_id"}, vectors, "embedding", entity.L2, topk, sp) if err != nil { log.Println("搜索相似失敗:", err) return } // 输出结果 for _, item := range resultx { ids = item.IDs.FieldData().GetScalars().GetLongData().GetData() } // 将结果缓存到 LRU 中 lruCache.Add(image.Id, ids) return ids } func main() { runtime.GOMAXPROCS(runtime.NumCPU() - 1) configFilePath := flag.String("config", "./data/config.yaml", "配置文件路径") flag.Parse() viper.SetConfigFile(*configFilePath) if err := viper.ReadInConfig(); err != nil { log.Println("读取配置文件失败", err) } config := viper.GetViper() models.InitConfig(config) models.ZincInit() api.InitDefault(config) mysqlConnection.Init() milvusConnection.Init() err := milvusConnection.Client.LoadCollection(context.Background(), "default", false) if err != nil { log.Println("Milvus load collection failed:", err) return } http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 w.Write([]byte("Hello World!")) }) if config.GetBool("oss.local") { fmt.Println("开启图像色调计算") go api.CheckColorNullRows(0) } if config.GetBool("gorse.open") { fmt.Println("开启用户偏好收集") api.GorseInit(config.GetString("gorse.host"), config.GetInt("gorse.port")) } schema, err := graphql.NewSchema(graphql.SchemaConfig{Query: graphql.NewObject(graphql.ObjectConfig{Name: "RootQuery", Fields: graphql.Fields{ "users": api.UserItems, "games": api.GameItems, "works": api.WorkItems, "collections": api.CollectionItems, "articles": api.ArticleItems, "article": api.ArticleItem, "images": api.ImageItems, "image": api.ImageItem, "searchs": api.SearchItems, }})}) if err != nil { log.Fatalf("failed to create new schema, error: %v", err) } http.Handle("/api", LogRequest(handler.New(&handler.Config{ Schema: &schema, Playground: true, Pretty: false, }))) // 获取图片信息列表(分页) http.HandleFunc("/api/images", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 // 获取查询条件(忽略空值) QueryConditions := func(key string) (list []string) { for _, item := range strings.Split(r.URL.Query().Get(key), ",") { if item != "" { list = append(list, fmt.Sprintf("'%s'", item)) } } return list } var conditions strings.Builder // 获取图片列表 var images ListView var image_list []Image images.Page, images.PageSize = stringToInt(r.URL.Query().Get("page"), 1), stringToInt(r.URL.Query().Get("pageSize"), 20) var ids []int64 if similar := QueryConditions("similar"); len(similar) > 0 { id, err := strconv.Atoi(strings.Trim(similar[0], "'")) if err != nil { log.Println("strconv.Atoi failed:", err) return } // 如果指定以某个图片为基准的相似图片列表范围, 获取相似图片ID的列表 ids = (&Image{Id: id}).GetSimilarImagesIdList("default") images.Total = len(ids) // 按照分页取相应的图片ID if len(ids) > images.Page*images.PageSize { ids = ids[(images.Page-1)*images.PageSize : images.Page*images.PageSize] } else { ids = ids[(images.Page-1)*images.PageSize:] } idsStr := make([]string, len(ids)) for i, v := range ids { idsStr[i] = strconv.FormatInt(v, 10) } if len(idsStr) > 0 { if conditions.Len() > 0 { conditions.WriteString(" AND") } else { conditions.WriteString(" WHERE") } conditions.WriteString(fmt.Sprintf(" id IN (%s)", strings.Join(idsStr, ","))) // 拼接查询条件 } } sql := fmt.Sprintf("SELECT id, width, height, content, update_time, create_time, user_id, article_id, article_category_top_id, praise_count, collect_count FROM web_images %s", conditions.String()) rows, err := mysqlConnection.Database.Query(sql) if err != nil { log.Println("获取图片列表失败", err) http.Error(w, err.Error(), http.StatusBadRequest) return } defer rows.Close() for rows.Next() { var image Image rows.Scan(&image.Id, &image.Width, &image.Height, &image.Content, &image.UpdateTime, &image.CreateTime, &image.User.ID, &image.Article.Id, &image.ArticleCategoryTopId, &image.PraiseCount, &image.CollectCount) image.UpdateTime = image.UpdateTime.UTC() image.CreateTime = image.CreateTime.UTC() image.Content = regexp.MustCompile(`http:`).ReplaceAllString(image.Content, "https:") image_list = append(image_list, image) } // 如果使用了相似图片列表范围, 需要按照相似图片ID原本的顺序重新排序, 需要注意的是, 相似图片ID列表中可能会包含不在当前页的图片ID if similar := QueryConditions("similar"); len(similar) > 0 { var image_list_sorted []Image for _, id := range ids { for _, image := range image_list { if image.Id == int(id) { image_list_sorted = append(image_list_sorted, image) } } } image_list = image_list_sorted } // 用户ID, 文章ID var user_ids []int var article_ids []int for _, image := range image_list { user_ids = append(user_ids, image.User.ID) article_ids = append(article_ids, image.Article.Id) } // 附加用户信息 users := models.QueryUserList(user_ids) for i, image := range image_list { for _, user := range users { if image.User.ID == user.ID { image_list[i].User = user } } } // 附加文章信息 articles := models.QueryArticleList(article_ids) for i, image := range image_list { for _, article := range articles { if image.Article.Id == article.Id { image_list[i].Article = article } } } // 将 []Image 转换为 []interface{} images.List = make([]interface{}, len(image_list)) for i, v := range image_list { images.List[i] = v } // 如果不是获取相似图像固定数量, 则从mysql获取总数 if similar := QueryConditions("similar"); len(similar) > 0 { // 固定数量 } else { // 获取总数 err = mysqlConnection.Database.QueryRow("SELECT COUNT(*) FROM web_images" + conditions.String()).Scan(&images.Total) if err != nil { log.Println("获取图片总数失败", err) http.Error(w, err.Error(), http.StatusBadRequest) return } } // 是否有下一页 images.Next = images.Total > images.Page*images.PageSize // 将对象转换为有缩进的JSON输出 data, err := json.MarshalIndent(images, "", " ") if err != nil { log.Println("转换图片列表失败", err) http.Error(w, err.Error(), http.StatusBadRequest) return } w.Header().Set("Content-Type", "application/json; charset=UTF-8") w.Write(data) }) // 获取标签列表 http.HandleFunc("/tags", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 // 标签的原理 // 1. 通过文章的 tag 字段, 获取所有的标签 // 2. 通过标签的 id, 获取标签的名称 // 热门权重指数的标签排序 // 1. 标签的权重指数 = (标签的文章数 * 标签的文章数) * 近期增幅 // 2. 标签的近期增幅 = (标签的文章数 - 标签的文章数) / 标签的文章数 // 标签是一个虚拟表, ORC 提取的数据都带有多个维度的比重概率(分布概率, 对比度概率, 文字大小, 文字重量, 词频概率, 词性概率, 词长概率, 词序概率) // 经过规则过滤后, 用动态调参的指数计算乘积作为权重, 权重仍达到某个阈值的数据才会被视为标签 // 获取查询条件(忽略空值), 超级简洁写法 QueryConditions := func(key string) (list []string) { if r.FormValue(key) != "" { list = strings.Split(r.FormValue(key), ",") } return } // 拼接查询条件, 超级简洁写法 conditions := "" if authors := QueryConditions("authors"); len(authors) > 0 { conditions += fmt.Sprintf(" AND author IN (%s)", strings.Join(authors, ",")) } if tags := QueryConditions("tags"); len(tags) > 0 { conditions += fmt.Sprintf(" AND tag IN (%s)", strings.Join(tags, ",")) } if categories := QueryConditions("categories"); len(categories) > 0 { conditions += fmt.Sprintf(" AND categorie IN (%s)", strings.Join(categories, ",")) } // 获取标签列表 var tags ListView tags.Page, tags.PageSize = stringToInt(r.FormValue("page"), 1), stringToInt(r.FormValue("pageSize"), 20) rows, err := mysqlConnection.Database.Query("SELECT id, name, update_time, create_time FROM web_tags"+conditions+" ORDER BY id DESC LIMIT ?, ?", (tags.Page-1)*tags.PageSize, tags.PageSize) if err != nil { log.Println(err) return } defer rows.Close() for rows.Next() { var tag Tag if err := rows.Scan(&tag.Id, &tag.Name, &tag.UpdateTime, &tag.CreateTime); err != nil { log.Println(err) continue } tags.List = append(tags.List, tag) } if err := rows.Err(); err != nil { log.Println(err) return } // 获取总数 if err := mysqlConnection.Database.QueryRow("SELECT COUNT(*) FROM web_tags" + conditions).Scan(&tags.Total); err != nil { log.Println(err) return } // 是否有下一页 tags.Next = tags.Total > tags.Page*tags.PageSize // 将对象转换为有缩进的JSON输出 json, err := json.MarshalIndent(tags, "", " ") if err != nil { log.Println(err) return } // 输出JSON w.Header().Set("Content-Type", "application/json") w.Write(json) }) // URL 格式: /img/{type}-{id}.{format}?width=320&height=320&fit=cover http.HandleFunc("/img/", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 reg := regexp.MustCompile(`^/img/([0-9a-zA-Z]+)-([0-9a-zA-Z]+).(jpg|jpeg|png|webp)$`) matches := reg.FindStringSubmatch(r.URL.Path) if len(matches) != 4 { http.Error(w, "URL 格式错误", http.StatusNotFound) return } group, id, format, width, height, fit := matches[1], matches[2], matches[3], stringToInt(r.URL.Query().Get("width"), 0), stringToInt(r.URL.Query().Get("height"), 0), r.URL.Query().Get("fit") content, err := mysqlConnection.GetImageContent(group, id) if err != nil { log.Println("获取图片失败", format, err) http.Error(w, err.Error(), http.StatusNotFound) return } var img models.Image if err := img.Init(content); err != nil { log.Println("初始化图片失败", format, err) http.Error(w, err.Error(), http.StatusBadRequest) return } data, err := img.ToWebP(width, height, fit) if err != nil { log.Println("转换图片失败", err) http.Error(w, err.Error(), http.StatusBadRequest) return } w.Header().Set("Content-Type", "image/webp") w.Header().Set("Cache-Control", "max-age=31536000") w.Write(data) }) // URL 格式: /webp/{type}-{id}-{version}-{width}-{height}-{fit}.{format} http.HandleFunc("/webp/", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 reg := regexp.MustCompile(`^/webp/([0-9a-zA-Z]+)-([0-9a-zA-Z]+)-([0-9a-zA-Z]+)-([0-9]+)-([0-9]+)-([a-zA-Z]+).(jpg|jpeg|png|webp)$`) matches := reg.FindStringSubmatch(r.URL.Path) if len(matches) != 8 { log.Println("URL 格式错误", matches) w.WriteHeader(http.StatusNotFound) return } group, id, version, width, height, fit, format := matches[1], matches[2], matches[3], stringToInt(matches[4], 0), stringToInt(matches[5], 0), matches[6], matches[7] content, err := mysqlConnection.GetImageContent(group, id) if err != nil { log.Println("获取图片失败", version, format, err) w.WriteHeader(http.StatusNotFound) return } var img models.Image if err := img.Init(content); err != nil { log.Println("初始化图片失败", version, format, err) w.WriteHeader(http.StatusNotFound) return } data, err := img.ToWebP(width, height, fit) if err != nil { log.Println("转换图片失败", err) w.WriteHeader(http.StatusBadRequest) return } w.Header().Set("Content-Type", "image/webp") w.Header().Set("Cache-Control", "max-age=31536000") w.Write(data) }) // 获取转换后的m3u8视频链接 http.HandleFunc("/video", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 queryParam := r.URL.Query().Get("url") safeParam, err := url.QueryUnescape(queryParam) if err != nil { log.Println("解码URL失败", err) w.WriteHeader(http.StatusBadRequest) return } fmt.Println("safeParam", safeParam) urls, err := models.GetVideoM3U8(safeParam) fmt.Println("urls", urls, err) if err != nil { log.Println("获取视频链接失败", err) w.WriteHeader(http.StatusBadRequest) return } // 将对象转换为有缩进的JSON输出 json, _ := json.MarshalIndent(urls, "", " ") w.Header().Set("Content-Type", "application/json") w.Write(json) }) // 从Viper中读取配置 port := viper.GetString("server.port") log.Println("Server is running at http://localhost:" + port) http.ListenAndServe(":"+port, nil) }