resume/internal/service/knowledge/vector_service.go

package knowledge

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"resume-platform/internal/model"
	"resume-platform/internal/repository"
	"sort"
)

type VectorService struct {
	embeddingService *EmbeddingService
	repo             repository.ResumeRepository
}

func NewVectorService(embeddingService *EmbeddingService, repo repository.ResumeRepository) *VectorService {
	return &VectorService{embeddingService: embeddingService, repo: repo}
}

func (s *VectorService) Store(ctx context.Context, content, sourceType, sourceID, sourceName, userID string) error {
	_, err := s.embeddingService.GenerateAndStore(ctx, content, sourceType, sourceID, sourceName, userID)
	return err
}

func (s *VectorService) StoreChunk(ctx context.Context, chunk *model.DocumentChunk, userID string, sourceName string) error {
	if chunk.Content == "" {
		return nil
	}

	embedding, err := s.embeddingService.Generate(ctx, chunk.Content)
	if err != nil {
		return err
	}

	embeddingJSON, _ := json.Marshal(embedding)
	chunk.Embedding = string(embeddingJSON)

	err = s.repo.CreateDocumentChunk(chunk)
	if err != nil {
		return err
	}

	contentHash := s.embeddingService.hashContent(chunk.Content)
	err = s.repo.CreateEmbedding(&model.Embedding{
		ID:          chunk.ID,
		ContentHash: contentHash,
		Embedding:   string(embeddingJSON),
		SourceType:  "document",
		SourceID:    chunk.DocumentID,
		SourceName:  sourceName,
		UserID:      userID,
	})

	return err
}

func (s *VectorService) Search(ctx context.Context, query string, topK int, userID string) ([]model.ChunkWithScore, error) {
	if query == "" {
		return nil, fmt.Errorf("empty query")
	}

	queryEmbedding, err := s.embeddingService.Generate(ctx, query)
	if err != nil {
		return nil, err
	}

	var allEmbeddings []*model.Embedding
	if userID != "" {
		allEmbeddings, err = s.repo.GetEmbeddingsByUserID(userID)
	} else {
		allEmbeddings, err = s.repo.GetAllEmbeddings()
	}
	if err != nil {
		return nil, err
	}

	var results []model.ChunkWithScore
	for _, emb := range allEmbeddings {
		if emb.Embedding == "" {
			continue
		}

		var embedding []float64
		if err := json.Unmarshal([]byte(emb.Embedding), &embedding); err != nil {
			continue
		}

		score := cosineSimilarity(queryEmbedding, embedding)
		if score > 0.3 {
			results = append(results, model.ChunkWithScore{
				Content:    emb.SourceName,
				Score:      score,
				SourceID:   emb.SourceID,
				SourceName: emb.SourceName,
				SourceType: emb.SourceType,
			})
		}
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	if len(results) > topK {
		results = results[:topK]
	}

	return results, nil
}

func (s *VectorService) SearchWithChunks(ctx context.Context, query string, topK int, userID string) ([]model.ChunkWithScore, error) {
	if query == "" {
		return nil, fmt.Errorf("empty query")
	}

	queryEmbedding, err := s.embeddingService.Generate(ctx, query)
	if err != nil {
		return nil, err
	}

	var allChunks []*model.DocumentChunk
	if userID != "" {
		allChunks, err = s.repo.GetDocumentChunksByUserID(userID)
	} else {
		allChunks, err = s.repo.GetAllDocumentChunks()
	}
	if err != nil {
		return nil, err
	}

	var results []model.ChunkWithScore
	for _, chunk := range allChunks {
		if chunk.Embedding == "" || chunk.Content == "" {
			continue
		}

		var embedding []float64
		if err := json.Unmarshal([]byte(chunk.Embedding), &embedding); err != nil {
			continue
		}

		score := cosineSimilarity(queryEmbedding, embedding)
		if score > 0.3 {
			doc, _ := s.repo.GetDocumentByID(chunk.DocumentID)
			sourceName := ""
			if doc != nil {
				sourceName = doc.Name
			}
			results = append(results, model.ChunkWithScore{
				Content:    chunk.Content,
				Score:      score,
				SourceID:   chunk.DocumentID,
				SourceName: sourceName,
				SourceType: "document",
			})
		}
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	if len(results) > topK {
		results = results[:topK]
	}

	return results, nil
}

func (s *VectorService) Delete(ctx context.Context, sourceType, sourceID string) error {
	return s.repo.DeleteEmbeddingsBySource(sourceType, sourceID)
}

func (s *VectorService) DeleteDocumentChunks(ctx context.Context, documentID string) error {
	_, err := s.repo.GetDocumentChunksByDocumentID(documentID)
	if err != nil {
		return err
	}
	return nil
}

func cosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) {
		return 0
	}

	var dotProduct, magA, magB float64
	for i := range a {
		dotProduct += a[i] * b[i]
		magA += a[i] * a[i]
		magB += b[i] * b[i]
	}

	if magA == 0 || magB == 0 {
		return 0
	}

	return dotProduct / (math.Sqrt(magA) * math.Sqrt(magB))
}