unified-media-manager/internal/service/semantic_search.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"time"
)

type SemanticSearchResult struct {
	ID        int64   `json:"id"`
	Title     string  `json:"title"`
	MediaType string  `json:"media_type"`
	Year      *int    `json:"year,omitempty"`
	Score     float64 `json:"score"`
	Overview  string  `json:"overview,omitempty"`
}

type ollamaEmbedRequest struct {
	Model  string `json:"model"`
	Prompt string `json:"prompt"`
}

type ollamaEmbedResponse struct {
	Embedding []float64 `json:"embedding"`
}

type qdrantSearchRequest struct {
	Vector      []float64 `json:"vector"`
	Limit       int       `json:"limit"`
	WithPayload bool      `json:"with_payload"`
}

type qdrantSearchResponse struct {
	Result []qdrantHit `json:"result"`
}

type qdrantHit struct {
	ID     string          `json:"id"`
	Score  float64         `json:"score"`
	Payload json.RawMessage `json:"payload"`
}

type qdrantPayload struct {
	ID        int64  `json:"id"`
	Title     string `json:"title"`
	MediaType string `json:"media_type"`
	Year      *int   `json:"year,omitempty"`
	Overview  string `json:"overview,omitempty"`
}

type SemanticSearchService struct {
	ollamaURL  string
	qdrantURL  string
	httpClient *http.Client
}

func NewSemanticSearchService(ollamaURL, qdrantURL string) *SemanticSearchService {
	return &SemanticSearchService{
		ollamaURL: ollamaURL,
		qdrantURL: qdrantURL,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

func (s *SemanticSearchService) Search(ctx context.Context, query string, k int) ([]SemanticSearchResult, error) {
	embedding, err := s.embed(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("generate embedding: %w", err)
	}

	results, err := s.searchQdrant(ctx, embedding, k)
	if err != nil {
		return nil, fmt.Errorf("search qdrant: %w", err)
	}

	return results, nil
}

func (s *SemanticSearchService) embed(ctx context.Context, text string) ([]float64, error) {
	reqBody := ollamaEmbedRequest{
		Model:  "nomic-embed-text",
		Prompt: text,
	}

	bodyBytes, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshal embed request: %w", err)
	}

	reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, s.ollamaURL+"/api/embeddings", bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, fmt.Errorf("create embed request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("ollama request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		slog.Error("ollama embed returned non-200", "status", resp.StatusCode, "body", string(respBody))
		return nil, fmt.Errorf("ollama returned status %d", resp.StatusCode)
	}

	var embedResp ollamaEmbedResponse
	if err := json.NewDecoder(resp.Body).Decode(&embedResp); err != nil {
		return nil, fmt.Errorf("decode ollama response: %w", err)
	}

	if len(embedResp.Embedding) == 0 {
		return nil, fmt.Errorf("ollama returned empty embedding")
	}

	return embedResp.Embedding, nil
}

func (s *SemanticSearchService) searchQdrant(ctx context.Context, embedding []float64, k int) ([]SemanticSearchResult, error) {
	reqBody := qdrantSearchRequest{
		Vector:      embedding,
		Limit:       k,
		WithPayload: true,
	}

	bodyBytes, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshal qdrant request: %w", err)
	}

	reqCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	url := fmt.Sprintf("%s/collections/media/points/search", s.qdrantURL)
	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, url, bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, fmt.Errorf("create qdrant request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("qdrant request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		slog.Error("qdrant search returned non-200", "status", resp.StatusCode, "body", string(respBody))
		return nil, fmt.Errorf("qdrant returned status %d", resp.StatusCode)
	}

	var searchResp qdrantSearchResponse
	if err := json.NewDecoder(resp.Body).Decode(&searchResp); err != nil {
		return nil, fmt.Errorf("decode qdrant response: %w", err)
	}

	results := make([]SemanticSearchResult, 0, len(searchResp.Result))
	for _, hit := range searchResp.Result {
		var payload qdrantPayload
		if err := json.Unmarshal(hit.Payload, &payload); err != nil {
			slog.Error("failed to unmarshal qdrant payload", "error", err)
			continue
		}

		results = append(results, SemanticSearchResult{
			ID:        payload.ID,
			Title:     payload.Title,
			MediaType: payload.MediaType,
			Year:      payload.Year,
			Score:     hit.Score,
			Overview:  payload.Overview,
		})
	}

	return results, nil
}