[pkgsite] internal, internal/frontend: add vector-search experiment flag
0 views
Skip to first unread message
Gopher Robot (Gerrit)
Aug 10, 2026, 5:58:27 PM (12 hours ago) Aug 10
to Ethan Lee, goph...@pubsubhelper.golang.org, golang-...@googlegroups.com, kokoro, golang...@luci-project-accounts.iam.gserviceaccount.com, Jonathan Amsterdam, golang-co...@googlegroups.com
Gopher Robot submitted the change with unreviewed changes![]( "Open in Gerrit")
Unreviewed changes
1 is the latest approved patch-set.
The change was submitted with unreviewed changes in the following files:
```
The name of the file: internal/frontend/search_test.go
Insertions: 0, Deletions: 1.
@@ -19,7 +19,6 @@
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/google/safehtml"
"golang.org/x/pkgsite/internal"
- "golang.org/x/pkgsite/internal/embeddings"
"golang.org/x/pkgsite/internal/experiment"
"golang.org/x/pkgsite/internal/fetchdatasource"
"golang.org/x/pkgsite/internal/frontend/page"
@@ -860,46 +859,57 @@
return s.FakeDataSource.Search(ctx, q, opts)
}
+type stubEmbedder struct {
+ vec []float32
+ err error
+}
+
+func (e *stubEmbedder) GenerateEmbeddings(ctx context.Context, texts []string, taskType string) ([][]float32, error) {
+ if e == nil {
+ return nil, nil
+ }
+ if e.err != nil {
+ return nil, e.err
+ }
+ return [][]float32{e.vec}, nil
+}
+
func TestFetchSearchPageWithVector(t *testing.T) {
ctx := experiment.NewContext(context.Background(), internal.ExperimentVectorSearch)
tests := []struct {
- name string
- httpStatus int
- httpResponse string
- wantVector []float32
+ name string
+ embedder VectorEmbedder
+ wantVector []float32
}{
{
- name: "embedding success",
- httpStatus: http.StatusOK,
- httpResponse: `{"predictions": [{"embeddings": {"values": [0.1, 0.2, 0.3]}}]}`,
- wantVector: []float32{0.1, 0.2, 0.3},
+ name: "embedding success",
+ embedder: &stubEmbedder{vec: []float32{0.1, 0.2, 0.3}},
+ wantVector: []float32{0.1, 0.2, 0.3},
},
{
- name: "embedding fail",
- httpStatus: http.StatusInternalServerError,
- httpResponse: `internal server error`,
- wantVector: nil,
+ name: "embedding fail",
+ embedder: &stubEmbedder{err: errors.New("embedding error")},
+ wantVector: nil,
+ },
+ {
+ name: "nil embedder interface",
+ embedder: nil,
+ wantVector: nil,
+ },
+ {
+ name: "typed nil embedder interface",
+ embedder: (*stubEmbedder)(nil),
+ wantVector: nil,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
- ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- w.WriteHeader(test.httpStatus)
- w.Header().Set("Content-Type", "application/json")
- fmt.Fprintln(w, test.httpResponse)
- }))
- defer ts.Close()
-
- mockClient := &embeddings.Client{
- HTTPClient: ts.Client(),
- BaseURL: ts.URL,
- }
ds := &spyDataSource{FakeDataSource: fakedatasource.New()}
- _, err := fetchSearchPage(ctx, ds, "http router", "", paginationParams{limit: 10, page: 1}, false, nil, mockClient)
+ _, err := fetchSearchPage(ctx, ds, "http router", "", paginationParams{limit: 10, page: 1}, false, nil, test.embedder)
if err != nil {
t.Fatalf("fetchSearchPage failed: %v", err)
}
```
```
The name of the file: internal/frontend/search.go
Insertions: 0, Deletions: 1.
@@ -21,7 +21,6 @@
"golang.org/x/mod/semver"
"golang.org/x/pkgsite/internal"
"golang.org/x/pkgsite/internal/derrors"
- "golang.org/x/pkgsite/internal/embeddings"
"golang.org/x/pkgsite/internal/experiment"
pagepkg "golang.org/x/pkgsite/internal/frontend/page"
"golang.org/x/pkgsite/internal/frontend/serrors"
@@ -61,7 +60,7 @@
page interface{ SetBasePage(pagepkg.BasePage) }
}
-func determineSearchAction(r *http.Request, ds internal.DataSource, vulnClient *vuln.Client, embeddingsClient *embeddings.Client) (*searchAction, error) {
+func determineSearchAction(r *http.Request, ds internal.DataSource, vulnClient *vuln.Client, embeddingsClient VectorEmbedder) (*searchAction, error) {
if r.Method != http.MethodGet && r.Method != http.MethodHead {
return nil, &serrors.ServerError{Status: http.StatusMethodNotAllowed}
}
@@ -245,14 +244,14 @@
// fetchSearchPage fetches data matching the search query from the database and
// returns a SearchPage.
func fetchSearchPage(ctx context.Context, ds internal.DataSource, cq, symbol string,
- pageParams paginationParams, searchSymbols bool, vulnClient *vuln.Client, embeddingsClient *embeddings.Client) (*SearchPage, error) {
+ pageParams paginationParams, searchSymbols bool, vulnClient *vuln.Client, embeddingsClient VectorEmbedder) (*SearchPage, error) {
maxResultCount := maxSearchOffset + pageParams.limit
var vec []float32
if embeddingsClient != nil && !searchSymbols && strings.TrimSpace(cq) != "" && experiment.IsActive(ctx, internal.ExperimentVectorSearch) {
embedCtx, cancel := context.WithTimeout(ctx, searchEmbeddingTimeout)
defer cancel()
- vecs, err := embeddingsClient.GenerateEmbeddings(embedCtx, []string{cq}, embeddings.TaskTypeQuery)
+ vecs, err := embeddingsClient.GenerateEmbeddings(embedCtx, []string{cq}, "RETRIEVAL_QUERY")
if err != nil {
log.Errorf(ctx, "failed to generate query vector for %q: %v", cq, err)
} else if len(vecs) > 0 {
```
Change information
Commit message:
internal, internal/frontend: add vector-search experiment flag
Add the vector-search experiment flag to internal/experiment.go and
gate query vector embedding generation in fetchSearchPage behind
experiment.IsActive(ctx, internal.ExperimentVectorSearch).
Change-Id: If0efb59aa05bbc1150e8970ff76f0928db704d7a
Reviewed-by: Jonathan Amsterdam <j...@google.com>
Auto-Submit: Ethan Lee <etha...@google.com>
LUCI-TryBot-Result: golang...@luci-project-accounts.iam.gserviceaccount.com <golang...@luci-project-accounts.iam.gserviceaccount.com>
kokoro-CI: kokoro <noreply...@google.com>
Files:
- M internal/experiment.go
- M internal/frontend/search.go
- M internal/frontend/search_test.go
Change size: XS
Delta: 3 files changed, 6 insertions(+), 2 deletions(-)
Branch: refs/heads/master
Submit Requirements:
Code-Review: +2 by Jonathan Amsterdam
TryBots-Pass: LUCI-TryBot-Result+1 by golang...@luci-project-accounts.iam.gserviceaccount.com
kokoro-CI-Passes: kokoro-CI+1 by kokoro
| Inspect html for hidden footers to help with email filtering. To unsubscribe visit settings. |
Reply all
Reply to author
Forward
0 new messages