清理旧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/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, 如果转换失败则返回默认值
|
||||
@@ -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 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)
|
||||
|
||||
|
Reference in New Issue
Block a user