diff --git a/bin/main.go b/bin/main.go index 0d2aa08..863d4b4 100644 --- a/bin/main.go +++ b/bin/main.go @@ -1,6 +1,7 @@ package main import ( + "context" "fmt" "log" "net/http" @@ -15,6 +16,7 @@ import ( "git.satori.love/gameui/webp/models" _ "github.com/go-sql-driver/mysql" + "github.com/milvus-io/milvus-sdk-go/v2/entity" ) // string 转换为 int, 如果转换失败则返回默认值 @@ -186,25 +188,70 @@ func main() { **/ }) - //// 获取相似图片列表 - //http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) { - // defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 - // id := "8888" - // // 先查询图片的向量在 mulvis 中是否存在 - // var collection_name = "default" // 图片集合名称 - // result, err := milvusConnection.Client.Query( - // context.Background(), // ctx - // collection_name, // CollectionName - // []string{}, // PartitionName - // fmt.Sprintf("id in [%s]", id), // expr - // []string{"book_id", "book_intro"}, // OutputFields - // ) - // if err != nil { - // log.Println(err) - // return - // } - // log.Println(result) - //}) + type Similar struct { + Id int64 `json:"id"` + ArticleId int64 `json:"article_id"` + Embedding []float32 `json:"embedding"` + } + + // 获取相似图片列表 + http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) { + defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 + id := "8888" + // 先查询图片的向量在 mulvis 中是否存在 + var collection_name = "default" // 图片集合名称 + result, err := milvusConnection.Client.Query( + context.Background(), // ctx + collection_name, // CollectionName + []string{}, // PartitionName + fmt.Sprintf("id in [%s]", id), // expr + []string{"id", "embedding", "article_id"}, // OutputFields + ) + if err != nil { + log.Println(err) + return + } + + // TODO: 不存在则重建向量 + var similar Similar + for _, item := range result { + if item.Name() == "id" { + similar.Id = item.FieldData().GetScalars().GetLongData().GetData()[0] + continue + } + if item.Name() == "article_id" { + similar.ArticleId = item.FieldData().GetScalars().GetLongData().GetData()[0] + continue + } + if item.Name() == "embedding" { + similar.Embedding = item.FieldData().GetVectors().GetFloatVector().Data + continue + } + } + + // 用向量查询相似图片 + sp, _ := entity.NewIndexFlatSearchParam() + vectors := []entity.Vector{ + entity.FloatVector(similar.Embedding), + } + resultx, err := milvusConnection.Client.Search( + context.Background(), // ctx + collection_name, // CollectionName + []string{}, // PartitionNames + "", // expr + []string{"id", "article_id"}, // OutputFields + vectors, // vectors + "embedding", // vectorField + entity.L2, // entity.MetricType + 10, // topK + sp, + ) + if err != nil { + log.Println(err) + return + } + println(resultx) + }) // 获取图片信息列表(分页) http.HandleFunc("/images", func(w http.ResponseWriter, r *http.Request) {