feat: 植物识别百科ai助手迁移
This commit is contained in:
Blizzard
@@ -1,24 +1,24 @@
|
||||
package logic
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
qdrant "github.com/qdrant/go-client/qdrant"
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
"google.golang.org/grpc/metadata"
|
||||
"gorm.io/gorm"
|
||||
|
||||
plantModel "sundynix-micro-go/app/plant/model"
|
||||
"sundynix-micro-go/app/plant/rpc/internal/config"
|
||||
)
|
||||
|
||||
func wikiVectorID(wikiID string) string {
|
||||
sum := md5.Sum([]byte("sundynix-plant-wiki:" + wikiID))
|
||||
return hex.EncodeToString(sum[:])
|
||||
return uuid.NewMD5(uuid.NameSpaceOID, []byte(wikiID)).String()
|
||||
}
|
||||
|
||||
func buildWikiVectorText(w plantModel.Wiki) string {
|
||||
@@ -30,131 +30,129 @@ func buildWikiVectorText(w plantModel.Wiki) string {
|
||||
w.FloweringShape, w.FlowerDiameter, w.Fruit)
|
||||
}
|
||||
|
||||
func embeddingModel(c config.Config) string {
|
||||
if c.Ai.EmbeddingModelName != "" {
|
||||
return c.Ai.EmbeddingModelName
|
||||
func getActiveAiConfig(db *gorm.DB) (*plantModel.SysAiConfig, error) {
|
||||
var cfg plantModel.SysAiConfig
|
||||
if err := db.Where("is_active = 1").First(&cfg).Error; err != nil {
|
||||
return nil, errors.New("数据库未找到已激活的 AI 配置")
|
||||
}
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
func embeddingModel(cfg *plantModel.SysAiConfig) string {
|
||||
if cfg.EmbeddingModelName != "" {
|
||||
return cfg.EmbeddingModelName
|
||||
}
|
||||
return "text-embedding-3-small"
|
||||
}
|
||||
|
||||
func createEmbedding(ctx context.Context, c config.Config, text string) ([]float32, error) {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"model": embeddingModel(c),
|
||||
"input": text,
|
||||
func createEmbedding(ctx context.Context, cfg *plantModel.SysAiConfig, text string) ([]float32, error) {
|
||||
config := openai.DefaultConfig(cfg.EmbeddingApiKey)
|
||||
if cfg.EmbeddingApiUrl != "" {
|
||||
config.BaseURL = cfg.EmbeddingApiUrl
|
||||
}
|
||||
client := openai.NewClientWithConfig(config)
|
||||
resp, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
|
||||
Input: []string{text},
|
||||
Model: openai.EmbeddingModel(embeddingModel(cfg)),
|
||||
})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.Ai.EmbeddingApiUrl, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+c.Ai.EmbeddingApiKey)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return nil, fmt.Errorf("embedding 请求失败: %s %s", resp.Status, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
var parsed struct {
|
||||
Data []struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(parsed.Data) == 0 || len(parsed.Data[0].Embedding) == 0 {
|
||||
if len(resp.Data) == 0 || len(resp.Data[0].Embedding) == 0 {
|
||||
return nil, errors.New("embedding 响应为空")
|
||||
}
|
||||
return parsed.Data[0].Embedding, nil
|
||||
return resp.Data[0].Embedding, nil
|
||||
}
|
||||
|
||||
func qdrantURL(c config.Config, path string) string {
|
||||
return strings.TrimRight(c.Ai.QdrantUrl, "/") + path
|
||||
}
|
||||
|
||||
func doQdrant(ctx context.Context, c config.Config, method, path string, body interface{}) error {
|
||||
var reader io.Reader
|
||||
if body != nil {
|
||||
raw, _ := json.Marshal(body)
|
||||
reader = bytes.NewReader(raw)
|
||||
func newQdrantConn(cfg *plantModel.SysAiConfig) (*grpc.ClientConn, context.Context, error) {
|
||||
addr := strings.TrimPrefix(cfg.QdrantUrl, "http://")
|
||||
addr = strings.TrimPrefix(addr, "https://")
|
||||
conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("qdrant grpc dial failed: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, qdrantURL(c, path), reader)
|
||||
ctx := context.Background()
|
||||
if cfg.QdrantApiKey != "" {
|
||||
ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("api-key", cfg.QdrantApiKey))
|
||||
}
|
||||
return conn, ctx, nil
|
||||
}
|
||||
|
||||
func ensureQdrantCollection(cfg *plantModel.SysAiConfig, dim int) error {
|
||||
conn, ctx, err := newQdrantConn(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.Ai.QdrantApiKey != "" {
|
||||
req.Header.Set("api-key", c.Ai.QdrantApiKey)
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return fmt.Errorf("qdrant 请求失败: %s %s", resp.Status, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func ensureQdrantCollection(ctx context.Context, c config.Config, dim int) error {
|
||||
getReq, err := http.NewRequestWithContext(ctx, http.MethodGet, qdrantURL(c, "/collections/"+c.Ai.QdrantCollection), nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if c.Ai.QdrantApiKey != "" {
|
||||
getReq.Header.Set("api-key", c.Ai.QdrantApiKey)
|
||||
}
|
||||
if resp, err := http.DefaultClient.Do(getReq); err == nil {
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
defer conn.Close()
|
||||
if dim <= 0 {
|
||||
dim = c.Ai.VectorDimension
|
||||
dim = cfg.VectorDimension
|
||||
}
|
||||
if dim <= 0 {
|
||||
dim = 1536
|
||||
}
|
||||
return doQdrant(ctx, c, http.MethodPut, "/collections/"+c.Ai.QdrantCollection, map[string]interface{}{
|
||||
"vectors": map[string]interface{}{
|
||||
"size": dim,
|
||||
"distance": "Cosine",
|
||||
collClient := qdrant.NewCollectionsClient(conn)
|
||||
if _, getErr := collClient.Get(ctx, &qdrant.GetCollectionInfoRequest{CollectionName: cfg.QdrantCollection}); getErr == nil {
|
||||
return nil
|
||||
}
|
||||
_, err = collClient.Create(ctx, &qdrant.CreateCollection{
|
||||
CollectionName: cfg.QdrantCollection,
|
||||
VectorsConfig: &qdrant.VectorsConfig{
|
||||
Config: &qdrant.VectorsConfig_Params{
|
||||
Params: &qdrant.VectorParams{Size: uint64(dim), Distance: qdrant.Distance_Cosine},
|
||||
},
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func upsertWikiVector(ctx context.Context, c config.Config, w plantModel.Wiki) error {
|
||||
func upsertWikiVector(ctx context.Context, cfg *plantModel.SysAiConfig, w plantModel.Wiki) error {
|
||||
text := buildWikiVectorText(w)
|
||||
vector, err := createEmbedding(ctx, c, text)
|
||||
vector, err := createEmbedding(ctx, cfg, text)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ensureQdrantCollection(ctx, c, len(vector)); err != nil {
|
||||
if err := ensureQdrantCollection(cfg, len(vector)); err != nil {
|
||||
return err
|
||||
}
|
||||
return doQdrant(ctx, c, http.MethodPut, "/collections/"+c.Ai.QdrantCollection+"/points?wait=true", map[string]interface{}{
|
||||
"points": []map[string]interface{}{
|
||||
{
|
||||
"id": wikiVectorID(w.ID),
|
||||
"vector": vector,
|
||||
"payload": map[string]interface{}{
|
||||
"wiki_id": w.ID,
|
||||
"name": w.Name,
|
||||
"full_text": text,
|
||||
conn, qdCtx, err := newQdrantConn(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
ptsClient := qdrant.NewPointsClient(conn)
|
||||
|
||||
_, err = ptsClient.Upsert(qdCtx, &qdrant.UpsertPoints{
|
||||
CollectionName: cfg.QdrantCollection,
|
||||
Points: []*qdrant.PointStruct{{
|
||||
Id: qdrant.NewID(wikiVectorID(w.ID)),
|
||||
Vectors: qdrant.NewVectors(vector...),
|
||||
Payload: map[string]*qdrant.Value{
|
||||
"wiki_id": qdrant.NewValueString(w.ID),
|
||||
"name": qdrant.NewValueString(w.Name),
|
||||
"full_text": qdrant.NewValueString(text),
|
||||
},
|
||||
}},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func deleteWikiVector(ctx context.Context, cfg *plantModel.SysAiConfig, wikiID string) error {
|
||||
conn, qdCtx, err := newQdrantConn(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
ptsClient := qdrant.NewPointsClient(conn)
|
||||
|
||||
_, err = ptsClient.Delete(qdCtx, &qdrant.DeletePoints{
|
||||
CollectionName: cfg.QdrantCollection,
|
||||
Points: &qdrant.PointsSelector{
|
||||
PointsSelectorOneOf: &qdrant.PointsSelector_Points{
|
||||
Points: &qdrant.PointsIdsList{
|
||||
Ids: []*qdrant.PointId{qdrant.NewID(wikiVectorID(wikiID))},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func deleteWikiVector(ctx context.Context, c config.Config, wikiID string) error {
|
||||
return doQdrant(ctx, c, http.MethodPost, "/collections/"+c.Ai.QdrantCollection+"/points/delete?wait=true", map[string]interface{}{
|
||||
"points": []string{wikiVectorID(wikiID)},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user