blob: 5783cebdff83cb3bec294de51a2c1bf4acd7a0a6 [file] [log] [blame]
// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package main
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"github.com/google/safehtml/template"
"golang.org/x/oscar/internal/llm"
"golang.org/x/oscar/internal/storage"
)
func (g *Gaby) handleSearch(w http.ResponseWriter, r *http.Request) {
query := r.FormValue("q")
if query == "" {
_, _ = io.WriteString(w, searchForm)
return
}
results, err := g.search(query)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
data, err := searchResultsHTML(query, results)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
_, _ = w.Write(data)
}
type searchResult struct {
Title string
VResult storage.VectorResult
}
func (g *Gaby) search(query string) ([]searchResult, error) {
vecs, err := g.embed.EmbedDocs(context.Background(), []llm.EmbedDoc{{Title: "", Text: query}})
if err != nil {
return nil, fmt.Errorf("EmbedDocs: %w", err)
}
vec := vecs[0]
var srs []searchResult
for _, r := range g.vector.Search(vec, 20) {
title := "?"
if d, ok := g.docs.Get(r.ID); ok {
title = d.Title
}
srs = append(srs, searchResult{title, r})
}
return srs, nil
}
func searchResultsHTML(query string, results []searchResult) ([]byte, error) {
data := struct {
Query string
Results []searchResult
}{
query, results,
}
var buf bytes.Buffer
if err := searchResultsTmpl.Execute(&buf, data); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
const searchForm = `
<!doctype html>
<html>
<head>
<title>Oscar Search</title>
</head>
<body>
<h1>Vector search</h1>
<form action="/search" method="GET">
<input type="text" name="q" required autofocus />
<input type="submit" value="Search"/>
</form
</body>
</html>
`
var searchResultsTmpl = template.Must(template.New("").Parse(`
<!doctype html>
<html>
<head>
<title>Oscar Search Results</title>
</head>
<body>
<h1>Search results for "{{.Query}}"</h1>
{{with .Results}}
{{range .Results}}
<p>{{.Title}}: {{.Vresult.ID}} ({{.VResult.Score}})</p>
{{end}}
{{else}}
No results.
{{end}}
</body>
</html>
`))