sift/internal/api/search.go

package api

import (
	"encoding/json"
	"net/http"
	"strings"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/search"
)

type v1SearchReq struct {
	Query              string   `json:"query"`
	Queries            []string `json:"queries"`
	Question           string   `json:"question"`
	K                  int      `json:"k"`
	Rerank             *bool    `json:"rerank"`
	Digest             bool     `json:"digest"`
	URL                string   `json:"url"`
	Limit              int      `json:"limit"`
	IncludeDomains     []string `json:"include_domains"`
	ExcludeDomains     []string `json:"exclude_domains"`
	IncludeDomainsMode string   `json:"include_domains_mode"`
	TimeRange          string   `json:"time_range"`
	StartDate          string   `json:"start_date"`
	EndDate            string   `json:"end_date"`
	Topic              string   `json:"topic"`
	ExactMatch         bool     `json:"exact_match"`
	MinScore           *float64 `json:"min_score"`
}

type v2SearchReq struct {
	Query  string `json:"query"`
	Limit  int    `json:"limit"`
	URL    string `json:"url"`
	Digest bool   `json:"digest"`
}

func (s *Server) v1search(w http.ResponseWriter, r *http.Request) {
	s.handleSearch(w, r, false)
}

func (s *Server) v2search(w http.ResponseWriter, r *http.Request) {
	s.handleSearch(w, r, true)
}

func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
	s.reqN.Add(1)
	c := s.caller(w, r)
	if c == nil {
		return
	}
	dec := json.NewDecoder(r.Body)
	query := ""
	var queries []string
	k := 0
	doRerank := true
	digest := false
	var f search.Filter
	if v2 {
		var req v2SearchReq
		if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
			writeErr(w, 400, "invalid query", "")
			return
		}
		query = req.Query
		k = req.Limit
		digest = req.Digest
		_ = req.URL
		if msg := search.ValidateQuery(query); msg != "" {
			writeErr(w, 400, msg, "")
			return
		}
		queries = []string{strings.TrimSpace(query)}
	} else {
		var req v1SearchReq
		if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
			writeErr(w, 400, "invalid query", "")
			return
		}
		k = req.K
		if k == 0 {
			k = req.Limit
		}
		if req.Rerank != nil {
			doRerank = *req.Rerank
		}
		digest = req.Digest
		_ = req.URL
		_ = req.Question
		raw := search.Filter{
			IncludeDomains:     req.IncludeDomains,
			ExcludeDomains:     req.ExcludeDomains,
			IncludeDomainsMode: req.IncludeDomainsMode,
			TimeRange:          req.TimeRange,
			StartDate:          req.StartDate,
			EndDate:            req.EndDate,
			Topic:              req.Topic,
			ExactMatch:         req.ExactMatch,
		}
		if req.MinScore != nil {
			raw.MinScore = *req.MinScore
		}
		var errMsg, hint string
		queries, f, errMsg, hint = search.NormalizeSearch(req.Query, req.Queries, raw)
		if errMsg != "" {
			writeErr(w, 400, errMsg, hint)
			return
		}
	}
	defK := search.DefaultK
	if v2 {
		defK = search.DefaultV2K
	}
	k = search.ClampK(k, defK)
	if s.Search == nil {
		writeErr(w, 422, "no search results", "")
		return
	}
	res := s.Search.RunFilter(r.Context(), queries, k, doRerank, c.Name, s.callerClass(c), f)
	warns := append([]string(nil), res.Warnings...)
	if digest {
		warns = append(warns, "digest ignored")
	}
	if !res.Cached && len(res.SourcesOK) == 0 && len(res.Hits) == 0 {
		writeErr(w, 422, "no search results", "")
		return
	}
	s.logSearch(c, queries, f, res)
	if v2 {
		web := make([]map[string]any, 0, len(res.Hits))
		for i, h := range res.Hits {
			web = append(web, map[string]any{
				"url":         h.URL,
				"title":       h.Title,
				"description": h.Snippet,
				"position":    i + 1,
			})
		}
		writeJSON(w, 200, map[string]any{"success": true, "data": map[string]any{"web": web}})
		return
	}
	results := make([]map[string]any, 0, len(res.Hits))
	for _, h := range res.Hits {
		row := map[string]any{
			"url":     h.URL,
			"title":   h.Title,
			"snippet": h.Snippet,
			"score":   h.Score,
			"source":  h.Source,
		}
		if h.PublishedDate != "" {
			row["published_date"] = h.PublishedDate
		}
		results = append(results, row)
	}
	writeJSON(w, 200, map[string]any{
		"results":  results,
		"warnings": warns,
		"sizes": map[string]any{
			"results":        len(results),
			"sources_ok":     res.SourcesOK,
			"sources_failed": res.SourcesFailed,
		},
		"timings": map[string]any{"search_ms": res.SearchMS, "rerank_ms": res.RerankMS},
		"cached":  res.Cached,
	})
}

func (s *Server) logSearch(c *auth.Caller, queries []string, f search.Filter, res search.RunResult) {
	qLen := search.QueryRuneLen(queries)
	cache := "miss"
	if res.Cached {
		cache = "hit"
	}
	src := strings.Join(res.Attempted, ",")
	if s.restrictedLog(c) {
		s.Log.Info("search", "caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache)
		return
	}
	kinds := search.FilterKinds(f)
	attrs := []any{"caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache, "rerank_ms", res.RerankMS, "n_queries", len(queries)}
	if kinds != "" {
		attrs = append(attrs, "filters", kinds)
	}
	s.Log.Info("search", attrs...)
}