ai-media-hub/backend/services/cse.go

package services

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	neturl "net/url"
	"strings"
	"time"
)

type SearchResult struct {
	Title        string `json:"title"`
	Link         string `json:"link"`
	DisplayLink  string `json:"displayLink"`
	Snippet      string `json:"snippet"`
	ThumbnailURL string `json:"thumbnailUrl"`
	Source       string `json:"source"`
}

type SearchService struct {
	APIKey        string
	ProjectID     string
	Location      string
	DataStoreID   string
	ServingConfig string
	Client        *http.Client
}

func NewSearchService(apiKey, projectID, location, dataStoreID, servingConfig string) *SearchService {
	if location == "" {
		location = "global"
	}
	if servingConfig == "" {
		servingConfig = "default_serving_config"
	}
	return &SearchService{
		APIKey:        apiKey,
		ProjectID:     projectID,
		Location:      location,
		DataStoreID:   dataStoreID,
		ServingConfig: servingConfig,
		Client:        &http.Client{Timeout: 20 * time.Second},
	}
}

func (s *SearchService) SearchMedia(query string) ([]SearchResult, error) {
	if s.APIKey == "" || s.ProjectID == "" || s.DataStoreID == "" {
		return nil, fmt.Errorf("vertex ai search credentials are not configured")
	}
	results, err := s.searchLite(query, true)
	if err != nil {
		results, err = s.searchLite(query, false)
		if err != nil {
			return nil, err
		}
	}
	return results, nil
}

func (s *SearchService) searchLite(query string, imageSearch bool) ([]SearchResult, error) {
	filteredQuery := strings.TrimSpace(query + " site:youtube.com OR site:tiktok.com OR site:envato.com OR site:artgrid.io")
	servingConfig := fmt.Sprintf(
		"projects/%s/locations/%s/dataStores/%s/servingConfigs/%s",
		s.ProjectID,
		s.Location,
		s.DataStoreID,
		s.ServingConfig,
	)

	params := map[string]any{
		"user_country_code": "us",
	}
	if imageSearch {
		params["searchType"] = 1
	}

	requestBody := map[string]any{
		"query":    filteredQuery,
		"pageSize": 25,
		"safeSearch": false,
		"languageCode": "ko-KR",
		"params": params,
		"contentSearchSpec": map[string]any{
			"snippetSpec": map[string]any{
				"returnSnippet": true,
			},
		},
	}

	body, _ := json.Marshal(requestBody)
	endpoint := fmt.Sprintf(
		"https://discoveryengine.googleapis.com/v1/%s:searchLite?key=%s",
		servingConfig,
		neturl.QueryEscape(s.APIKey),
	)
	resp, err := s.Client.Post(endpoint, "application/json", strings.NewReader(string(body)))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return nil, fmt.Errorf("vertex ai search returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(data)))
	}

	var payload struct {
		Results []struct {
			Document struct {
				StructData        map[string]any `json:"structData"`
				DerivedStructData map[string]any `json:"derivedStructData"`
			} `json:"document"`
		} `json:"results"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return nil, err
	}

	results := make([]SearchResult, 0, len(payload.Results))
	for _, item := range payload.Results {
		link := firstString(item.Document.StructData, "link", "url", "uri")
		title := firstString(item.Document.StructData, "title", "name")
		displayLink := firstString(item.Document.StructData, "site_name", "displayLink")
		snippet := firstString(item.Document.DerivedStructData, "snippets", "snippet")
		thumb := firstString(item.Document.DerivedStructData, "link", "thumbnail", "image", "image_url")
		if thumb == "" {
			thumb = firstString(item.Document.StructData, "thumbnail", "image", "image_url")
		}
		if thumb == "" || link == "" {
			continue
		}
		results = append(results, SearchResult{
			Title:        title,
			Link:         link,
			DisplayLink:  displayLink,
			Snippet:      snippet,
			ThumbnailURL: thumb,
			Source:       inferSource(displayLink + " " + link),
		})
	}
	return results, nil
}

func firstString(values map[string]any, keys ...string) string {
	for _, key := range keys {
		value, ok := values[key]
		if !ok {
			continue
		}
		switch typed := value.(type) {
		case string:
			if typed != "" {
				return typed
			}
		case []any:
			for _, item := range typed {
				if text, ok := item.(string); ok && text != "" {
					return text
				}
				if mapped, ok := item.(map[string]any); ok {
					if text := firstString(mapped, "snippet", "htmlSnippet", "url"); text != "" {
						return text
					}
				}
			}
		case map[string]any:
			if text := firstString(typed, "snippet", "htmlSnippet", "url"); text != "" {
				return text
			}
		}
	}
	return ""
}

func inferSource(displayLink string) string {
	switch {
	case strings.Contains(displayLink, "youtube"):
		return "YouTube"
	case strings.Contains(displayLink, "tiktok"):
		return "TikTok"
	case strings.Contains(displayLink, "envato"):
		return "Envato"
	case strings.Contains(displayLink, "artgrid"):
		return "Artgrid"
	default:
		return displayLink
	}
}