From 5b53cc6e11dbcfa4b5296e78008cbc3311e572b6 Mon Sep 17 00:00:00 2001 From: AI Assistant Date: Thu, 12 Mar 2026 16:31:45 +0900 Subject: [PATCH] Migrate search to Vertex AI and enhance preview modal --- backend/handlers/api.go | 1 + backend/main.go | 8 +- backend/services/cse.go | 199 +++++++++++++++++++++------------------- frontend/app.js | 15 +++ frontend/index.html | 15 ++- unraid-template.xml | 7 +- worker/downloader.py | 27 +++++- 7 files changed, 170 insertions(+), 102 deletions(-) diff --git a/backend/handlers/api.go b/backend/handlers/api.go index 1220118..884d467 100644 --- a/backend/handlers/api.go +++ b/backend/handlers/api.go @@ -67,6 +67,7 @@ func (h *Hub) Remove(conn *websocket.Conn) { type PreviewResponse struct { Title string `json:"title"` Thumbnail string `json:"thumbnail"` + PreviewStreamURL string `json:"previewStreamUrl"` Duration string `json:"duration"` DurationSeconds int `json:"durationSeconds"` StartDefault string `json:"startDefault"` diff --git a/backend/main.go b/backend/main.go index 08da603..e686caa 100644 --- a/backend/main.go +++ b/backend/main.go @@ -34,7 +34,13 @@ func main() { DB: db, DownloadsDir: downloadsDir, WorkerScript: workerScript, - SearchService: services.NewSearchService(os.Getenv("GOOGLE_CSE_API_KEY"), os.Getenv("GOOGLE_CSE_CX")), + SearchService: services.NewSearchService( + os.Getenv("VERTEX_AI_SEARCH_API_KEY"), + os.Getenv("VERTEX_AI_SEARCH_PROJECT_ID"), + os.Getenv("VERTEX_AI_SEARCH_LOCATION"), + os.Getenv("VERTEX_AI_SEARCH_DATA_STORE_ID"), + os.Getenv("VERTEX_AI_SEARCH_SERVING_CONFIG"), + ), GeminiService: services.NewGeminiService(os.Getenv("GEMINI_API_KEY")), Hub: handlers.NewHub(), } diff --git a/backend/services/cse.go b/backend/services/cse.go index b218f17..cff9b96 100644 --- a/backend/services/cse.go +++ b/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 "" diff --git a/frontend/app.js b/frontend/app.js index 1a58a01..ba07dba 100644 --- a/frontend/app.js +++ b/frontend/app.js @@ -15,6 +15,7 @@ const downloadResult = document.getElementById("downloadResult"); const cardTemplate = document.getElementById("searchCardTemplate"); const previewModal = document.getElementById("previewModal"); const previewTitle = document.getElementById("previewTitle"); +const previewVideo = document.getElementById("previewVideo"); const previewThumbnail = document.getElementById("previewThumbnail"); const previewDuration = document.getElementById("previewDuration"); const qualitySelect = document.getElementById("qualitySelect"); @@ -117,6 +118,17 @@ function openPreviewModal(preview) { previewTitle.textContent = preview.title; previewThumbnail.src = preview.thumbnail; previewThumbnail.alt = preview.title; + previewVideo.pause(); + previewVideo.removeAttribute("src"); + previewVideo.load(); + if (preview.previewStreamUrl) { + previewVideo.src = preview.previewStreamUrl; + previewVideo.classList.remove("hidden"); + previewThumbnail.classList.add("hidden"); + } else { + previewVideo.classList.add("hidden"); + previewThumbnail.classList.remove("hidden"); + } previewDuration.textContent = preview.duration; qualitySelect.innerHTML = ""; for (const item of preview.qualities || []) { @@ -132,6 +144,9 @@ function openPreviewModal(preview) { } function closeModal() { + previewVideo.pause(); + previewVideo.removeAttribute("src"); + previewVideo.load(); previewModal.classList.add("hidden"); previewModal.classList.remove("flex"); pendingDownload = null; diff --git a/frontend/index.html b/frontend/index.html index 75356e1..1e9fa0e 100644 --- a/frontend/index.html +++ b/frontend/index.html @@ -60,10 +60,6 @@

Direct Downloader & Crop

-
- - -

@@ -83,6 +79,7 @@
+
@@ -92,6 +89,16 @@
+
+ + +