lpra/internal/server/search.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"html"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Web search for the assistant, when the user turned «Поиск» on: DuckDuckGo's
// plain pages, no key and nothing to install. The light page first, the
// HTML one if that fails; DuckDuckGo turns away too many requests from one
// address, and then the assistant is told so.

type searchResult struct {
	Title   string `json:"title"`
	URL     string `json:"url"`
	Snippet string `json:"snippet"`
}

var (
	reLiteLink    = regexp.MustCompile(`(?s)<a[^>]+href="([^"]+)"[^>]*class='result-link'[^>]*>(.*?)</a>`)
	reLiteSnippet = regexp.MustCompile(`(?s)<td class='result-snippet'>(.*?)</td>`)
	reHTMLLink    = regexp.MustCompile(`(?s)<a[^>]+class="result__a"[^>]+href="([^"]+)"[^>]*>(.*?)</a>`)
	reHTMLSnippet = regexp.MustCompile(`(?s)<a[^>]+class="result__snippet"[^>]*>(.*?)</a>`)
	reTags        = regexp.MustCompile(`<[^>]+>`)
)

// text is a bit of a results page as plain text.
func text(s string) string {
	return strings.Join(strings.Fields(html.UnescapeString(reTags.ReplaceAllString(s, ""))), " ")
}

// realURL takes DuckDuckGo's redirect (/l/?uddg=…) back to the address.
func realURL(h string) string {
	h = html.UnescapeString(h)
	if strings.HasPrefix(h, "//") {
		h = "https:" + h
	}
	if u, err := url.Parse(h); err == nil {
		if t := u.Query().Get("uddg"); t != "" {
			return t
		}
	}
	return h
}

// parseDDG reads the results of the light or the HTML page.
func parseDDG(page string) []searchResult {
	links, snips := reLiteLink.FindAllStringSubmatch(page, -1), reLiteSnippet.FindAllStringSubmatch(page, -1)
	if len(links) == 0 {
		links, snips = reHTMLLink.FindAllStringSubmatch(page, -1), reHTMLSnippet.FindAllStringSubmatch(page, -1)
	}
	var out []searchResult
	for i, l := range links {
		r := searchResult{Title: text(l[2]), URL: realURL(l[1])}
		if i < len(snips) {
			r.Snippet = text(snips[i][1])
		}
		// Ads lead back through DuckDuckGo itself: left out.
		if strings.Contains(r.URL, "duckduckgo.com/y.js") || r.Title == "" {
			continue
		}
		out = append(out, r)
	}
	return out
}

var searchClient = &http.Client{Timeout: 15 * time.Second}

// The engine is the admins' choice (Настройки › ИИ-ассистент): DuckDuckGo,
// a SearXNG of one's own, Brave Search or Tavily; when the chosen one fails,
// DuckDuckGo answers instead.
const (
	cfgSearchEngine = "ai_search"
	cfgSearchURL    = "ai_search_url"
	cfgSearchKey    = "ai_search_key"
)

type searchConf struct{ Engine, URL, Key string }

func (s *Server) searchConf() searchConf {
	e, _ := s.db.GetConfig(cfgSearchEngine)
	u, _ := s.db.GetConfig(cfgSearchURL)
	k, _ := s.db.GetConfig(cfgSearchKey)
	return searchConf{Engine: e, URL: strings.TrimRight(u, "/"), Key: k}
}

var (
	braveURL  = "https://api.search.brave.com/res/v1/web/search"
	tavilyURL = "https://api.tavily.com/search"
	ddgSearch = webSearch
)

// searchWith asks the engine c names; DuckDuckGo if that one fails.
func searchWith(ctx context.Context, c searchConf, q string, n int) ([]searchResult, error) {
	var rs []searchResult
	var err error
	switch c.Engine {
	case "searxng":
		rs, err = searchSearx(ctx, c, q)
	case "brave":
		rs, err = searchBrave(ctx, c, q, n)
	case "tavily":
		rs, err = searchTavily(ctx, c, q, n)
	default:
		return ddgSearch(ctx, q, n)
	}
	if err != nil || len(rs) == 0 {
		return ddgSearch(ctx, q, n)
	}
	if len(rs) > n {
		rs = rs[:n]
	}
	return rs, nil
}

// getJSON fetches url (with the headers) into v.
func getJSON(ctx context.Context, method, url string, body io.Reader, hdr map[string]string, v any) error {
	req, err := http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	for k, val := range hdr {
		req.Header.Set(k, val)
	}
	resp, err := searchClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return errors.New("поисковик ответил " + resp.Status)
	}
	return json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(v)
}

// searchSearx: a SearXNG with JSON output on (search.formats: json).
func searchSearx(ctx context.Context, c searchConf, q string) ([]searchResult, error) {
	if c.URL == "" {
		return nil, errors.New("не задан адрес SearXNG")
	}
	var out struct {
		Results []struct{ Title, URL, Content string } `json:"results"`
	}
	if err := getJSON(ctx, http.MethodGet, c.URL+"/search?format=json&q="+url.QueryEscape(q), nil, nil, &out); err != nil {
		return nil, err
	}
	var rs []searchResult
	for _, r := range out.Results {
		rs = append(rs, searchResult{Title: r.Title, URL: r.URL, Snippet: text(r.Content)})
	}
	return rs, nil
}

func searchBrave(ctx context.Context, c searchConf, q string, n int) ([]searchResult, error) {
	var out struct {
		Web struct {
			Results []struct{ Title, URL, Description string } `json:"results"`
		} `json:"web"`
	}
	err := getJSON(ctx, http.MethodGet, braveURL+"?count="+strconv.Itoa(n)+"&q="+url.QueryEscape(q), nil,
		map[string]string{"X-Subscription-Token": c.Key}, &out)
	if err != nil {
		return nil, err
	}
	var rs []searchResult
	for _, r := range out.Web.Results {
		rs = append(rs, searchResult{Title: text(r.Title), URL: r.URL, Snippet: text(r.Description)})
	}
	return rs, nil
}

func searchTavily(ctx context.Context, c searchConf, q string, n int) ([]searchResult, error) {
	b, _ := json.Marshal(map[string]any{"query": q, "max_results": n})
	var out struct {
		Results []struct{ Title, URL, Content string } `json:"results"`
	}
	err := getJSON(ctx, http.MethodPost, tavilyURL, strings.NewReader(string(b)),
		map[string]string{"Authorization": "Bearer " + c.Key, "Content-Type": "application/json"}, &out)
	if err != nil {
		return nil, err
	}
	var rs []searchResult
	for _, r := range out.Results {
		rs = append(rs, searchResult{Title: r.Title, URL: r.URL, Snippet: r.Content})
	}
	return rs, nil
}

// webSearch asks DuckDuckGo for q: up to n results.
func webSearch(ctx context.Context, q string, n int) ([]searchResult, error) {
	var last error
	for _, base := range []string{"https://lite.duckduckgo.com/lite/?q=", "https://html.duckduckgo.com/html/?q="} {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, base+url.QueryEscape(q), nil)
		req.Header.Set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0 Safari/537.36")
		req.Header.Set("Accept-Language", "ru,en;q=0.8")
		resp, err := searchClient.Do(req)
		if err != nil {
			last = err
			continue
		}
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
		resp.Body.Close()
		rs := parseDDG(string(b))
		if resp.StatusCode == http.StatusOK && len(rs) > 0 {
			if len(rs) > n {
				rs = rs[:n]
			}
			return rs, nil
		}
		if resp.StatusCode == http.StatusOK {
			last = errors.New("ничего не нашлось")
		} else {
			last = errors.New("поисковик не ответил (" + resp.Status + "), возможно, временно ограничил запросы")
		}
	}
	return nil, last
}

// POST /api/panel/ai/search - {q}: the results, for a user with the
// assistant on (the «Поиск» switch is the page's: it lets the model ask).
func (s *Server) handleAISearch(w http.ResponseWriter, r *http.Request) {
	a := s.currentAdmin(r)
	if on, _ := s.db.GetConfig(aiUserKey(a.ID)); on != "1" {
		writeJSON(w, map[string]any{"ok": false, "error": "ассистент выключен в профиле"})
		return
	}
	var in struct {
		Q string `json:"q"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<10)).Decode(&in); err != nil || strings.TrimSpace(in.Q) == "" {
		writeJSON(w, map[string]any{"ok": false, "error": "пустой запрос"})
		return
	}
	rs, err := searchWith(r.Context(), s.searchConf(), strings.TrimSpace(in.Q), 6)
	if err != nil {
		writeJSON(w, map[string]any{"ok": false, "error": err.Error()})
		return
	}
	writeJSON(w, map[string]any{"ok": true, "data": rs})
}