Реализована операции Milvus для управления документами и встраиванием, включая функции вставки, запроса и удаления. Внедрите архитектуру RAG с LLM и сервисами встраивания. Добавьте обработку текста для фрагментации и конкатенации. Создайте автономный скрипт для настройки и управления Milvus. Разработайте комплексные тесты API для обработки документов и взаимодействия с LLM, включая имитации для сервисов. Расширьте возможности конфигурации пользователя с помощью дополнительных настроек YAML.

2025-09-19 11:38:31 +03:00
parent 8e7aab5181
commit 636096fd34
38 changed files with 3420 additions and 28 deletions
--- a/internal/database/database_milvus.go
+++ b/internal/database/database_milvus.go
@@ -0,0 +1,142 @@
+package database
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"sort"
+
+	"easy_rag/internal/models"
+	"easy_rag/internal/pkg/database/milvus"
+)
+
+// implement database interface for milvus
+type Milvus struct {
+	Host   string
+	Client *milvus.Client
+}
+
+func NewMilvus(host string) *Milvus {
+
+	milviusClient, err := milvus.NewClient(host)
+	if err != nil {
+		panic(err)
+	}
+
+	return &Milvus{
+		Host:   host,
+		Client: milviusClient,
+	}
+}
+
+func (m *Milvus) SaveDocument(document models.Document) error {
+	// for now lets use context background
+	ctx := context.Background()
+
+	return m.Client.InsertDocuments(ctx, []models.Document{document})
+}
+
+func (m *Milvus) SaveEmbeddings(embeddings []models.Embedding) error {
+	ctx := context.Background()
+	return m.Client.InsertEmbeddings(ctx, embeddings)
+}
+
+func (m *Milvus) GetDocumentInfo(id string) (models.Document, error) {
+	ctx := context.Background()
+	doc, err := m.Client.GetDocumentByID(ctx, id)
+
+	if err != nil {
+		return models.Document{}, err
+	}
+
+	if len(doc) == 0 {
+		return models.Document{}, nil
+	}
+	return models.Document{
+		ID:             doc["ID"].(string),
+		Link:           doc["Link"].(string),
+		Filename:       doc["Filename"].(string),
+		Category:       doc["Category"].(string),
+		EmbeddingModel: doc["EmbeddingModel"].(string),
+		Summary:        doc["Summary"].(string),
+		Metadata:       doc["Metadata"].(map[string]string),
+	}, nil
+}
+
+func (m *Milvus) GetDocument(id string) (models.Document, error) {
+	ctx := context.Background()
+	doc, err := m.Client.GetDocumentByID(ctx, id)
+	if err != nil {
+		return models.Document{}, err
+	}
+
+	embeds, err := m.Client.GetAllEmbeddingByDocID(ctx, id)
+
+	if err != nil {
+		return models.Document{}, err
+	}
+
+	// order embed by order
+	sort.Slice(embeds, func(i, j int) bool {
+		return embeds[i].Order < embeds[j].Order
+	})
+
+	// concatenate text chunks
+	var buf bytes.Buffer
+	for _, embed := range embeds {
+		buf.WriteString(embed.TextChunk)
+	}
+
+	textChunks := buf.String()
+
+	if len(doc) == 0 {
+		return models.Document{}, nil
+	}
+	return models.Document{
+		ID:             doc["ID"].(string),
+		Content:        textChunks,
+		Link:           doc["Link"].(string),
+		Filename:       doc["Filename"].(string),
+		Category:       doc["Category"].(string),
+		EmbeddingModel: doc["EmbeddingModel"].(string),
+		Summary:        doc["Summary"].(string),
+		Metadata:       doc["Metadata"].(map[string]string),
+	}, nil
+}
+
+func (m *Milvus) Search(vector [][]float32) ([]models.Embedding, error) {
+	ctx := context.Background()
+	results, err := m.Client.Search(ctx, vector, 10)
+	if err != nil {
+		return nil, err
+	}
+
+	return results, nil
+}
+
+func (m *Milvus) ListDocuments() ([]models.Document, error) {
+	ctx := context.Background()
+
+	docs, err := m.Client.GetAllDocuments(ctx)
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to get docs: %w", err)
+	}
+
+	return docs, nil
+}
+
+func (m *Milvus) DeleteDocument(id string) error {
+	ctx := context.Background()
+	err := m.Client.DeleteDocument(ctx, id)
+	if err != nil {
+		return err
+	}
+
+	err = m.Client.DeleteEmbedding(ctx, id)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}