Реализована операции 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/pkg/database/milvus/client.go
+++ b/internal/pkg/database/milvus/client.go
@@ -0,0 +1,169 @@
+package milvus
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+)
+
+type Client struct {
+	Instance client.Client
+}
+
+// InitMilvusClient initializes the Milvus client and returns a wrapper around it.
+func NewClient(milvusAddr string) (*Client, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	c, err := client.NewClient(ctx, client.Config{Address: milvusAddr})
+	if err != nil {
+		log.Printf("Failed to connect to Milvus: %v", err)
+		return nil, err
+	}
+
+	client := &Client{Instance: c}
+
+	err = client.EnsureCollections(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return client, nil
+}
+
+// EnsureCollections ensures that the required collections ("documents" and "chunks") exist.
+// If they don't exist, it creates them based on the predefined structs.
+func (m *Client) EnsureCollections(ctx context.Context) error {
+	collections := []struct {
+		Name       string
+		Schema     *entity.Schema
+		IndexField string
+		IndexType  string
+		MetricType entity.MetricType
+		Nlist      int
+	}{
+		{
+			Name:       "documents",
+			Schema:     createDocumentSchema(),
+			IndexField: "Vector", // Indexing the Vector field for similarity search
+			IndexType:  "IVF_FLAT",
+			MetricType: entity.L2,
+			Nlist:      10, // Number of clusters for IVF_FLAT index
+		},
+		{
+			Name:       "chunks",
+			Schema:     createEmbeddingSchema(),
+			IndexField: "Vector", // Indexing the Vector field for similarity search
+			IndexType:  "IVF_FLAT",
+			MetricType: entity.L2,
+			Nlist:      10,
+		},
+	}
+
+	for _, collection := range collections {
+		// drop collection
+		// err := m.Instance.DropCollection(ctx, collection.Name)
+		// if err != nil {
+		// 	return fmt.Errorf("failed to drop collection '%s': %w", collection.Name, err)
+		// }
+
+		// Ensure the collection exists
+		exists, err := m.Instance.HasCollection(ctx, collection.Name)
+		if err != nil {
+			return fmt.Errorf("failed to check collection existence: %w", err)
+		}
+
+		if !exists {
+			err := m.Instance.CreateCollection(ctx, collection.Schema, entity.DefaultShardNumber)
+			if err != nil {
+				return fmt.Errorf("failed to create collection '%s': %w", collection.Name, err)
+			}
+			log.Printf("Collection '%s' created successfully", collection.Name)
+		} else {
+			log.Printf("Collection '%s' already exists", collection.Name)
+		}
+
+		// Ensure the default partition exists
+		hasPartition, err := m.Instance.HasPartition(ctx, collection.Name, "_default")
+		if err != nil {
+			return fmt.Errorf("failed to check default partition for collection '%s': %w", collection.Name, err)
+		}
+
+		if !hasPartition {
+			err = m.Instance.CreatePartition(ctx, collection.Name, "_default")
+			if err != nil {
+				return fmt.Errorf("failed to create default partition for collection '%s': %w", collection.Name, err)
+			}
+			log.Printf("Default partition created for collection '%s'", collection.Name)
+		}
+
+		// Skip index creation if IndexField is empty
+		if collection.IndexField == "" {
+			continue
+		}
+
+		// Ensure the index exists
+		log.Printf("Creating index on field '%s' for collection '%s'", collection.IndexField, collection.Name)
+
+		idx, err := entity.NewIndexIvfFlat(collection.MetricType, collection.Nlist)
+		if err != nil {
+			return fmt.Errorf("failed to create IVF_FLAT index: %w", err)
+		}
+
+		err = m.Instance.CreateIndex(ctx, collection.Name, collection.IndexField, idx, false)
+		if err != nil {
+			return fmt.Errorf("failed to create index on field '%s' for collection '%s': %w", collection.IndexField, collection.Name, err)
+		}
+
+		log.Printf("Index created on field '%s' for collection '%s'", collection.IndexField, collection.Name)
+	}
+
+	err := m.Instance.LoadCollection(ctx, "documents", false)
+	if err != nil {
+		log.Fatalf("failed to load collection, err: %v", err)
+	}
+
+	err = m.Instance.LoadCollection(ctx, "chunks", false)
+	if err != nil {
+		log.Fatalf("failed to load collection, err: %v", err)
+	}
+
+	return nil
+}
+
+// Helper functions for creating schemas
+func createDocumentSchema() *entity.Schema {
+	return entity.NewSchema().
+		WithName("documents").
+		WithDescription("Collection for storing documents").
+		WithField(entity.NewField().WithName("ID").WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(512)).
+		WithField(entity.NewField().WithName("Content").WithDataType(entity.FieldTypeVarChar).WithMaxLength(65535)).
+		WithField(entity.NewField().WithName("Link").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512)).
+		WithField(entity.NewField().WithName("Filename").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512)).
+		WithField(entity.NewField().WithName("Category").WithDataType(entity.FieldTypeVarChar).WithMaxLength(8048)).
+		WithField(entity.NewField().WithName("EmbeddingModel").WithDataType(entity.FieldTypeVarChar).WithMaxLength(256)).
+		WithField(entity.NewField().WithName("Summary").WithDataType(entity.FieldTypeVarChar).WithMaxLength(65535)).
+		WithField(entity.NewField().WithName("Metadata").WithDataType(entity.FieldTypeVarChar).WithMaxLength(65535)).
+		WithField(entity.NewField().WithName("Vector").WithDataType(entity.FieldTypeFloatVector).WithDim(1024)) // bge-m3
+}
+
+func createEmbeddingSchema() *entity.Schema {
+	return entity.NewSchema().
+		WithName("chunks").
+		WithDescription("Collection for storing document embeddings").
+		WithField(entity.NewField().WithName("ID").WithDataType(entity.FieldTypeVarChar).WithIsPrimaryKey(true).WithMaxLength(512)).
+		WithField(entity.NewField().WithName("DocumentID").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512)).
+		WithField(entity.NewField().WithName("Vector").WithDataType(entity.FieldTypeFloatVector).WithDim(1024)). // bge-m3
+		WithField(entity.NewField().WithName("TextChunk").WithDataType(entity.FieldTypeVarChar).WithMaxLength(65535)).
+		WithField(entity.NewField().WithName("Dimension").WithDataType(entity.FieldTypeInt32)).
+		WithField(entity.NewField().WithName("Order").WithDataType(entity.FieldTypeInt32))
+}
+
+// Close closes the Milvus client connection.
+func (m *Client) Close() {
+	m.Instance.Close()
+}
--- a/internal/pkg/database/milvus/client_test.go
+++ b/internal/pkg/database/milvus/client_test.go
@@ -0,0 +1,32 @@
+package milvus
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestNewClient(t *testing.T) {
+	type args struct {
+		milvusAddr string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    *Client
+		wantErr bool
+	}{
+		// TODO: Add test cases.
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := NewClient(tt.args.milvusAddr)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewClient() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("NewClient() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
--- a/internal/pkg/database/milvus/helpers.go
+++ b/internal/pkg/database/milvus/helpers.go
@@ -0,0 +1,276 @@
+package milvus
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"easy_rag/internal/models"
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+)
+
+// Helper functions for extracting data
+func extractIDs(docs []models.Document) []string {
+	ids := make([]string, len(docs))
+	for i, doc := range docs {
+		ids[i] = doc.ID
+	}
+	return ids
+}
+
+// extractLinks extracts the "Link" field from the documents.
+func extractLinks(docs []models.Document) []string {
+	links := make([]string, len(docs))
+	for i, doc := range docs {
+		links[i] = doc.Link
+	}
+	return links
+}
+
+// extractFilenames extracts the "Filename" field from the documents.
+func extractFilenames(docs []models.Document) []string {
+	filenames := make([]string, len(docs))
+	for i, doc := range docs {
+		filenames[i] = doc.Filename
+	}
+	return filenames
+}
+
+// extractCategories extracts the "Category" field from the documents as a comma-separated string.
+func extractCategories(docs []models.Document) []string {
+	categories := make([]string, len(docs))
+	for i, doc := range docs {
+		categories[i] = fmt.Sprintf("%v", doc.Category)
+	}
+	return categories
+}
+
+// extractEmbeddingModels extracts the "EmbeddingModel" field from the documents.
+func extractEmbeddingModels(docs []models.Document) []string {
+	models := make([]string, len(docs))
+	for i, doc := range docs {
+		models[i] = doc.EmbeddingModel
+	}
+	return models
+}
+
+// extractSummaries extracts the "Summary" field from the documents.
+func extractSummaries(docs []models.Document) []string {
+	summaries := make([]string, len(docs))
+	for i, doc := range docs {
+		summaries[i] = doc.Summary
+	}
+	return summaries
+}
+
+// extractMetadata extracts the "Metadata" field from the documents as a JSON string.
+func extractMetadata(docs []models.Document) []string {
+	metadata := make([]string, len(docs))
+	for i, doc := range docs {
+		metaBytes, _ := json.Marshal(doc.Metadata)
+		metadata[i] = string(metaBytes)
+	}
+	return metadata
+}
+
+func convertToMetadata(metadata string) map[string]string {
+	var metadataMap map[string]string
+	json.Unmarshal([]byte(metadata), &metadataMap)
+	return metadataMap
+}
+
+func extractContents(docs []models.Document) []string {
+	contents := make([]string, len(docs))
+	for i, doc := range docs {
+		contents[i] = doc.Content
+	}
+	return contents
+}
+
+// extractEmbeddingIDs extracts the "ID" field from the embeddings.
+func extractEmbeddingIDs(embeddings []models.Embedding) []string {
+	ids := make([]string, len(embeddings))
+	for i, embedding := range embeddings {
+		ids[i] = embedding.ID
+	}
+	return ids
+}
+
+// extractDocumentIDs extracts the "DocumentID" field from the embeddings.
+func extractDocumentIDs(embeddings []models.Embedding) []string {
+	documentIDs := make([]string, len(embeddings))
+	for i, embedding := range embeddings {
+		documentIDs[i] = embedding.DocumentID
+	}
+	return documentIDs
+}
+
+// extractVectors extracts the "Vector" field from the embeddings.
+func extractVectors(embeddings []models.Embedding) [][]float32 {
+	vectors := make([][]float32, len(embeddings))
+	for i, embedding := range embeddings {
+		vectors[i] = embedding.Vector // Direct assignment since it's already []float32
+	}
+	return vectors
+}
+
+// extractVectorsDocs extracts the "Vector" field from the documents.
+func extractVectorsDocs(docs []models.Document) [][]float32 {
+	vectors := make([][]float32, len(docs))
+	for i, doc := range docs {
+		vectors[i] = doc.Vector // Direct assignment since it's already []float32
+	}
+	return vectors
+}
+
+// extractTextChunks extracts the "TextChunk" field from the embeddings.
+func extractTextChunks(embeddings []models.Embedding) []string {
+	textChunks := make([]string, len(embeddings))
+	for i, embedding := range embeddings {
+		textChunks[i] = embedding.TextChunk
+	}
+	return textChunks
+}
+
+// extractDimensions extracts the "Dimension" field from the embeddings.
+func extractDimensions(embeddings []models.Embedding) []int32 {
+	dimensions := make([]int32, len(embeddings))
+	for i, embedding := range embeddings {
+		dimensions[i] = int32(embedding.Dimension)
+	}
+	return dimensions
+}
+
+// extractOrders extracts the "Order" field from the embeddings.
+func extractOrders(embeddings []models.Embedding) []int32 {
+	orders := make([]int32, len(embeddings))
+	for i, embedding := range embeddings {
+		orders[i] = int32(embedding.Order)
+	}
+	return orders
+}
+
+func transformResultSet(rs client.ResultSet, outputFields ...string) ([]map[string]interface{}, error) {
+	if rs == nil || rs.Len() == 0 {
+		return nil, fmt.Errorf("empty result set")
+	}
+
+	results := []map[string]interface{}{}
+
+	for i := 0; i < rs.Len(); i++ { // Iterate through rows
+		row := map[string]interface{}{}
+
+		for _, fieldName := range outputFields {
+			column := rs.GetColumn(fieldName)
+			if column == nil {
+				return nil, fmt.Errorf("column %s does not exist in result set", fieldName)
+			}
+
+			switch column.Type() {
+			case entity.FieldTypeInt64:
+				value, err := column.GetAsInt64(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting int64 value for column %s, row %d: %w", fieldName, i, err)
+				}
+				row[fieldName] = value
+
+			case entity.FieldTypeInt32:
+				value, err := column.GetAsInt64(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting int64 value for column %s, row %d: %w", fieldName, i, err)
+				}
+				row[fieldName] = value
+
+			case entity.FieldTypeFloat:
+				value, err := column.GetAsDouble(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting float value for column %s, row %d: %w", fieldName, i, err)
+				}
+				row[fieldName] = value
+
+			case entity.FieldTypeDouble:
+				value, err := column.GetAsDouble(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting double value for column %s, row %d: %w", fieldName, i, err)
+				}
+				row[fieldName] = value
+
+			case entity.FieldTypeVarChar:
+				value, err := column.GetAsString(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting string value for column %s, row %d: %w", fieldName, i, err)
+				}
+				row[fieldName] = value
+
+			default:
+				return nil, fmt.Errorf("unsupported field type for column %s", fieldName)
+			}
+		}
+
+		results = append(results, row)
+	}
+
+	return results, nil
+}
+
+func transformSearchResultSet(rs client.SearchResult, outputFields ...string) ([]map[string]interface{}, error) {
+	if rs.ResultCount == 0 {
+		return nil, fmt.Errorf("empty result set")
+	}
+
+	result := make([]map[string]interface{}, rs.ResultCount)
+
+	for i := 0; i < rs.ResultCount; i++ { // Iterate through rows
+		result[i] = make(map[string]interface{})
+		for _, fieldName := range outputFields {
+			column := rs.Fields.GetColumn(fieldName)
+			result[i]["Score"] = rs.Scores[i]
+
+			if column == nil {
+				return nil, fmt.Errorf("column %s does not exist in result set", fieldName)
+			}
+
+			switch column.Type() {
+			case entity.FieldTypeInt64:
+				value, err := column.GetAsInt64(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting int64 value for column %s, row %d: %w", fieldName, i, err)
+				}
+				result[i][fieldName] = value
+
+			case entity.FieldTypeInt32:
+				value, err := column.GetAsInt64(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting int64 value for column %s, row %d: %w", fieldName, i, err)
+				}
+				result[i][fieldName] = value
+
+			case entity.FieldTypeFloat:
+				value, err := column.GetAsDouble(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting float value for column %s, row %d: %w", fieldName, i, err)
+				}
+				result[i][fieldName] = value
+
+			case entity.FieldTypeDouble:
+				value, err := column.GetAsDouble(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting double value for column %s, row %d: %w", fieldName, i, err)
+				}
+				result[i][fieldName] = value
+
+			case entity.FieldTypeVarChar:
+				value, err := column.GetAsString(i)
+				if err != nil {
+					return nil, fmt.Errorf("error getting string value for column %s, row %d: %w", fieldName, i, err)
+				}
+				result[i][fieldName] = value
+
+			default:
+				return nil, fmt.Errorf("unsupported field type for column %s", fieldName)
+			}
+		}
+	}
+
+	return result, nil
+}
--- a/internal/pkg/database/milvus/operations.go
+++ b/internal/pkg/database/milvus/operations.go
@@ -0,0 +1,270 @@
+package milvus
+
+import (
+	"context"
+	"fmt"
+	"sort"
+
+	"easy_rag/internal/models"
+
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+)
+
+// InsertDocuments inserts documents into the "documents" collection.
+func (m *Client) InsertDocuments(ctx context.Context, docs []models.Document) error {
+	idColumn := entity.NewColumnVarChar("ID", extractIDs(docs))
+	contentColumn := entity.NewColumnVarChar("Content", extractContents(docs))
+	linkColumn := entity.NewColumnVarChar("Link", extractLinks(docs))
+	filenameColumn := entity.NewColumnVarChar("Filename", extractFilenames(docs))
+	categoryColumn := entity.NewColumnVarChar("Category", extractCategories(docs))
+	embeddingModelColumn := entity.NewColumnVarChar("EmbeddingModel", extractEmbeddingModels(docs))
+	summaryColumn := entity.NewColumnVarChar("Summary", extractSummaries(docs))
+	metadataColumn := entity.NewColumnVarChar("Metadata", extractMetadata(docs))
+	vectorColumn := entity.NewColumnFloatVector("Vector", 1024, extractVectorsDocs(docs))
+	// Insert the data
+	_, err := m.Instance.Insert(ctx, "documents", "_default", idColumn, contentColumn, linkColumn, filenameColumn,
+		categoryColumn, embeddingModelColumn, summaryColumn, metadataColumn, vectorColumn)
+	if err != nil {
+		return fmt.Errorf("failed to insert documents: %w", err)
+	}
+
+	// Flush the collection
+	err = m.Instance.Flush(ctx, "documents", false)
+	if err != nil {
+		return fmt.Errorf("failed to flush documents collection: %w", err)
+	}
+
+	return nil
+}
+
+// InsertEmbeddings inserts embeddings into the "chunks" collection.
+func (m *Client) InsertEmbeddings(ctx context.Context, embeddings []models.Embedding) error {
+	idColumn := entity.NewColumnVarChar("ID", extractEmbeddingIDs(embeddings))
+	documentIDColumn := entity.NewColumnVarChar("DocumentID", extractDocumentIDs(embeddings))
+	vectorColumn := entity.NewColumnFloatVector("Vector", 1024, extractVectors(embeddings))
+	textChunkColumn := entity.NewColumnVarChar("TextChunk", extractTextChunks(embeddings))
+	dimensionColumn := entity.NewColumnInt32("Dimension", extractDimensions(embeddings))
+	orderColumn := entity.NewColumnInt32("Order", extractOrders(embeddings))
+
+	_, err := m.Instance.Insert(ctx, "chunks", "_default", idColumn, documentIDColumn, vectorColumn,
+		textChunkColumn, dimensionColumn, orderColumn)
+
+	if err != nil {
+		return fmt.Errorf("failed to insert embeddings: %w", err)
+	}
+
+	err = m.Instance.Flush(ctx, "chunks", false)
+	if err != nil {
+		return fmt.Errorf("failed to flush chunks collection: %w", err)
+	}
+
+	return nil
+}
+
+// GetDocumentByID retrieves a document from the "documents" collection by ID.
+func (m *Client) GetDocumentByID(ctx context.Context, id string) (map[string]interface{}, error) {
+	collectionName := "documents"
+	expr := fmt.Sprintf("ID == '%s'", id)
+	projections := []string{"ID", "Content", "Link", "Filename", "Category", "EmbeddingModel", "Summary", "Metadata"} // Fetch all fields
+
+	results, err := m.Instance.Query(ctx, collectionName, nil, expr, projections)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query document by ID: %w", err)
+	}
+
+	if len(results) == 0 {
+		return nil, fmt.Errorf("document with ID '%s' not found", id)
+	}
+
+	mp, err := transformResultSet(results, "ID", "Content", "Link", "Filename", "Category", "EmbeddingModel", "Summary", "Metadata")
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal document: %w", err)
+	}
+
+	// convert metadata to map
+	mp[0]["Metadata"] = convertToMetadata(mp[0]["Metadata"].(string))
+
+	return mp[0], err
+}
+
+// GetAllDocuments retrieves all documents from the "documents" collection.
+func (m *Client) GetAllDocuments(ctx context.Context) ([]models.Document, error) {
+	collectionName := "documents"
+	projections := []string{"ID", "Content", "Link", "Filename", "Category", "EmbeddingModel", "Summary", "Metadata"} // Fetch all fields
+	expr := ""
+
+	rs, err := m.Instance.Query(ctx, collectionName, nil, expr, projections, client.WithLimit(1000))
+	if err != nil {
+		return nil, fmt.Errorf("failed to query all documents: %w", err)
+	}
+
+	if len(rs) == 0 {
+		return nil, fmt.Errorf("no documents found in the collection")
+	}
+
+	results, err := transformResultSet(rs, "ID", "Content", "Link", "Filename", "Category", "EmbeddingModel", "Summary", "Metadata")
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal all documents: %w", err)
+	}
+
+	var docs []models.Document = make([]models.Document, len(results))
+	for i, result := range results {
+		docs[i] = models.Document{
+			ID:             result["ID"].(string),
+			Content:        result["Content"].(string),
+			Link:           result["Link"].(string),
+			Filename:       result["Filename"].(string),
+			Category:       result["Category"].(string),
+			EmbeddingModel: result["EmbeddingModel"].(string),
+			Summary:        result["Summary"].(string),
+			Metadata:       convertToMetadata(results[0]["Metadata"].(string)),
+		}
+	}
+
+	return docs, nil
+}
+
+// GetAllEmbeddingByDocID retrieves all embeddings linked to a specific DocumentID from the "chunks" collection.
+func (m *Client) GetAllEmbeddingByDocID(ctx context.Context, documentID string) ([]models.Embedding, error) {
+	collectionName := "chunks"
+	projections := []string{"ID", "DocumentID", "TextChunk", "Order"} // Fetch all fields
+	expr := fmt.Sprintf("DocumentID == '%s'", documentID)
+
+	rs, err := m.Instance.Query(ctx, collectionName, nil, expr, projections, client.WithLimit(1000))
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to query embeddings by DocumentID: %w", err)
+	}
+
+	if rs.Len() == 0 {
+		return nil, fmt.Errorf("no embeddings found for DocumentID '%s'", documentID)
+	}
+
+	results, err := transformResultSet(rs, "ID", "DocumentID", "TextChunk", "Order")
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal all documents: %w", err)
+	}
+
+	var embeddings []models.Embedding = make([]models.Embedding, rs.Len())
+
+	for i, result := range results {
+		embeddings[i] = models.Embedding{
+			ID:         result["ID"].(string),
+			DocumentID: result["DocumentID"].(string),
+			TextChunk:  result["TextChunk"].(string),
+			Order:      result["Order"].(int64),
+		}
+	}
+
+	return embeddings, nil
+}
+
+func (m *Client) Search(ctx context.Context, vectors [][]float32, topK int) ([]models.Embedding, error) {
+	const (
+		collectionName = "chunks"
+		vectorDim      = 1024 // Replace with your actual vector dimension
+	)
+	projections := []string{"ID", "DocumentID", "TextChunk", "Order"}
+	metricType := entity.L2 // Default metric type
+
+	// Validate and convert input vectors
+	searchVectors, err := validateAndConvertVectors(vectors, vectorDim)
+	if err != nil {
+		return nil, err
+	}
+
+	// Set search parameters
+	searchParams, err := entity.NewIndexIvfFlatSearchParam(16) // 16 is the number of clusters for IVF_FLAT index
+	if err != nil {
+		return nil, fmt.Errorf("failed to create search params: %w", err)
+	}
+
+	// Perform the search
+	searchResults, err := m.Instance.Search(ctx, collectionName, nil, "", projections, searchVectors, "Vector", metricType, topK, searchParams, client.WithLimit(10))
+	if err != nil {
+		return nil, fmt.Errorf("failed to search collection: %w", err)
+	}
+
+	// Process search results
+	embeddings, err := processSearchResults(searchResults)
+	if err != nil {
+		return nil, fmt.Errorf("failed to process search results: %w", err)
+	}
+
+	return embeddings, nil
+}
+
+// validateAndConvertVectors validates vector dimensions and converts them to Milvus-compatible format.
+func validateAndConvertVectors(vectors [][]float32, expectedDim int) ([]entity.Vector, error) {
+	searchVectors := make([]entity.Vector, len(vectors))
+	for i, vector := range vectors {
+		if len(vector) != expectedDim {
+			return nil, fmt.Errorf("vector dimension mismatch: expected %d, got %d", expectedDim, len(vector))
+		}
+		searchVectors[i] = entity.FloatVector(vector)
+	}
+	return searchVectors, nil
+}
+
+// processSearchResults transforms and aggregates the search results into embeddings and sorts by score.
+func processSearchResults(results []client.SearchResult) ([]models.Embedding, error) {
+	var embeddings []models.Embedding
+
+	for _, result := range results {
+		for i := 0; i < result.ResultCount; i++ {
+			embeddingMap, err := transformSearchResultSet(result, "ID", "DocumentID", "TextChunk", "Order")
+			if err != nil {
+				return nil, fmt.Errorf("failed to transform search result set: %w", err)
+			}
+
+			for _, embedding := range embeddingMap {
+				embeddings = append(embeddings, models.Embedding{
+					ID:         embedding["ID"].(string),
+					DocumentID: embedding["DocumentID"].(string),
+					TextChunk:  embedding["TextChunk"].(string),
+					Order:      embedding["Order"].(int64), // Assuming 'Order' is a float64 type
+					Score:      embedding["Score"].(float32),
+				})
+			}
+		}
+	}
+
+	// Sort embeddings by score in descending order (higher is better)
+	sort.Slice(embeddings, func(i, j int) bool {
+		return embeddings[i].Score > embeddings[j].Score
+	})
+
+	return embeddings, nil
+}
+
+// DeleteDocument deletes a document from the "documents" collection by ID.
+func (m *Client) DeleteDocument(ctx context.Context, id string) error {
+	collectionName := "documents"
+	partitionName := "_default"
+	expr := fmt.Sprintf("ID == '%s'", id)
+
+	err := m.Instance.Delete(ctx, collectionName, partitionName, expr)
+	if err != nil {
+		return fmt.Errorf("failed to delete document by ID: %w", err)
+	}
+
+	return nil
+}
+
+// DeleteEmbedding deletes an embedding from the "chunks" collection by ID.
+func (m *Client) DeleteEmbedding(ctx context.Context, id string) error {
+	collectionName := "chunks"
+	partitionName := "_default"
+	expr := fmt.Sprintf("DocumentID == '%s'", id)
+
+	err := m.Instance.Delete(ctx, collectionName, partitionName, expr)
+	if err != nil {
+		return fmt.Errorf("failed to delete embedding by DocumentID: %w", err)
+	}
+
+	return nil
+}