Migrate search to Vertex AI and enhance preview modal

backend/services/cse.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/url"
+	neturl "net/url"
 	"strings"
 	"time"
 )
@@ -20,71 +20,82 @@ type SearchResult struct {
 }
 
 type SearchService struct {
-	APIKey string
-	CX     string
-	Client *http.Client
+	APIKey        string
+	ProjectID     string
+	Location      string
+	DataStoreID   string
+	ServingConfig string
+	Client        *http.Client
 }
 
-func NewSearchService(apiKey, cx string) *SearchService {
+func NewSearchService(apiKey, projectID, location, dataStoreID, servingConfig string) *SearchService {
+	if location == "" {
+		location = "global"
+	}
+	if servingConfig == "" {
+		servingConfig = "default_serving_config"
+	}
 	return &SearchService{
-		APIKey: apiKey,
-		CX:     cx,
-		Client: &http.Client{Timeout: 20 * time.Second},
+		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.CX == "" {
-		return nil, fmt.Errorf("google cse credentials are not configured")
+	if s.APIKey == "" || s.ProjectID == "" || s.DataStoreID == "" {
+		return nil, fmt.Errorf("vertex ai search credentials are not configured")
 	}
-
-	domains := []string{"youtube.com", "tiktok.com", "envato.com", "artgrid.io"}
-	siteQuery := strings.Join(domains, " OR site:")
-	fullQuery := fmt.Sprintf("%s (site:%s)", query, siteQuery)
-
-	values := url.Values{}
-	values.Set("key", s.APIKey)
-	values.Set("cx", s.CX)
-	values.Set("q", fullQuery)
-	values.Set("num", "10")
-	values.Set("safe", "off")
-
-	results := make([]SearchResult, 0, 30)
-	seen := map[string]bool{}
-	for _, start := range []string{"1", "11", "21"} {
-		pageResults, err := s.fetchPage(values, start, true)
+	results, err := s.searchLite(query, true)
+	if err != nil {
+		results, err = s.searchLite(query, false)
 		if err != nil {
-			pageResults, err = s.fetchPage(values, start, false)
-			if err != nil {
-				return nil, err
-			}
-		}
-
-		for _, item := range pageResults {
-			if item.Link == "" || item.ThumbnailURL == "" || seen[item.Link] {
-				continue
-			}
-			seen[item.Link] = true
-			results = append(results, item)
+			return nil, err
 		}
 	}
 	return results, nil
 }
 
-func (s *SearchService) fetchPage(values url.Values, start string, imageSearch bool) ([]SearchResult, error) {
-	pageValues := url.Values{}
-	for key, items := range values {
-		for _, item := range items {
-			pageValues.Add(key, item)
-		}
+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",
 	}
-	pageValues.Set("start", start)
 	if imageSearch {
-		pageValues.Set("searchType", "image")
+		params["searchType"] = 1
 	}
 
-	endpoint := "https://www.googleapis.com/customsearch/v1?" + pageValues.Encode()
-	resp, err := s.Client.Get(endpoint)
+	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
 	}
@@ -92,69 +103,73 @@ func (s *SearchService) fetchPage(values url.Values, start string, imageSearch b
 
 	if resp.StatusCode >= 300 {
 		data, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
-		return nil, fmt.Errorf("google cse returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(data)))
+		return nil, fmt.Errorf("vertex ai search returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(data)))
 	}
 
 	var payload struct {
-		Items []struct {
-			Title       string `json:"title"`
-			Link        string `json:"link"`
-			DisplayLink string `json:"displayLink"`
-			Snippet     string `json:"snippet"`
-			Image       struct {
-				ThumbnailLink string `json:"thumbnailLink"`
-			} `json:"image"`
-			Pagemap struct {
-				CSEImage []struct {
-					Src string `json:"src"`
-				} `json:"cse_image"`
-				CSEThumbnail []struct {
-					Src string `json:"src"`
-				} `json:"cse_thumbnail"`
-				Metatags []map[string]string `json:"metatags"`
-			} `json:"pagemap"`
-		} `json:"items"`
+		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.Items))
-	for _, item := range payload.Items {
-		thumb := item.Image.ThumbnailLink
+	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 = extractThumbnail(item.Pagemap)
+			thumb = firstString(item.Document.StructData, "thumbnail", "image", "image_url")
+		}
+		if thumb == "" || link == "" {
+			continue
 		}
 		results = append(results, SearchResult{
-			Title:        item.Title,
-			Link:         item.Link,
-			DisplayLink:  item.DisplayLink,
-			Snippet:      item.Snippet,
+			Title:        title,
+			Link:         link,
+			DisplayLink:  displayLink,
+			Snippet:      snippet,
 			ThumbnailURL: thumb,
-			Source:       inferSource(item.DisplayLink),
+			Source:       inferSource(displayLink + " " + link),
 		})
 	}
 	return results, nil
 }
 
-func extractThumbnail(pagemap struct {
-	CSEImage     []struct{ Src string "json:\"src\"" }            "json:\"cse_image\""
-	CSEThumbnail []struct{ Src string "json:\"src\"" }            "json:\"cse_thumbnail\""
-	Metatags     []map[string]string                              "json:\"metatags\""
-}) string {
-	if len(pagemap.CSEThumbnail) > 0 && pagemap.CSEThumbnail[0].Src != "" {
-		return pagemap.CSEThumbnail[0].Src
-	}
-	if len(pagemap.CSEImage) > 0 && pagemap.CSEImage[0].Src != "" {
-		return pagemap.CSEImage[0].Src
-	}
-	for _, tag := range pagemap.Metatags {
-		if value := tag["og:image"]; value != "" {
-			return value
+func firstString(values map[string]any, keys ...string) string {
+	for _, key := range keys {
+		value, ok := values[key]
+		if !ok {
+			continue
 		}
-		if value := tag["twitter:image"]; value != "" {
-			return value
+		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 ""