清理旧API
This commit is contained in:
125
bin/main.go
125
bin/main.go
@@ -22,10 +22,7 @@ import (
|
|||||||
_ "github.com/go-sql-driver/mysql"
|
_ "github.com/go-sql-driver/mysql"
|
||||||
"github.com/graphql-go/graphql"
|
"github.com/graphql-go/graphql"
|
||||||
"github.com/graphql-go/handler"
|
"github.com/graphql-go/handler"
|
||||||
"github.com/milvus-io/milvus-sdk-go/v2/entity"
|
|
||||||
"github.com/spf13/viper"
|
"github.com/spf13/viper"
|
||||||
|
|
||||||
lru "github.com/hashicorp/golang-lru/v2"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
// string 转换为 int, 如果转换失败则返回默认值
|
// string 转换为 int, 如果转换失败则返回默认值
|
||||||
@@ -75,131 +72,9 @@ func LogRequest(next http.Handler) http.Handler {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
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 mysqlConnection models.MysqlConnection
|
||||||
var milvusConnection models.MilvusConnection
|
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() {
|
func main() {
|
||||||
runtime.GOMAXPROCS(runtime.NumCPU() - 1)
|
runtime.GOMAXPROCS(runtime.NumCPU() - 1)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user