Gather additional statistics at index time (#64426)

Adds a new `GetCodyContextAlternatives` GraphQL resolver that returns a
list of `ContextList`, each of which is a distinct list of results that
match the user query, along with a name that describes how that result
list was computed. The idea here is to be able to show alternative
context fetched with different methods/rankers, which can be displayed
by the client to help us iterate on context quality and ordering.

The existing context fetching mechanism is preserved and given the
`keyword($query)` name.

We add a new experimental context list, which modifies the user query to
do term expansion using a table of keywords. This keyword table is
computed in the `RepositoryReindex` GraphQL endpoint and stored in
Redis. We do not worry about clean up, because this feature is
experimental.

All of this is protected by the feature flag `enhanced-index`. Existing
production behavior should not be modified.

---------

Co-authored-by: Rishabh Mehrotra <[email protected]>

Author: beyang

.gitignore

@@ -76,8 +76,6 @@ eb-bundle.zip
 .bin/
 .env
 
-/vendor/
-
 .gtm/
 
 cmd/gitserver/debug

cmd/frontend/graphqlbackend/BUILD.bazel

@@ -246,10 +246,12 @@ go_library(
         "//cmd/frontend/internal/backend",
         "//cmd/frontend/internal/cloneurls",
         "//cmd/frontend/internal/cody",
+        "//cmd/frontend/internal/codycontext",
         "//cmd/frontend/internal/conf/validation",
         "//cmd/frontend/internal/highlight",
         "//cmd/frontend/internal/inventory",
         "//cmd/frontend/internal/processrestart",
+        "//cmd/frontend/internal/search/idf",
         "//cmd/frontend/internal/search/logs",
         "//cmd/frontend/internal/ssc",
         "//cmd/frontend/internal/suspiciousnames",
@@ -374,6 +376,7 @@ go_library(
         "@com_github_masterminds_semver//:semver",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_prometheus_client_golang//prometheus/promauto",
+        "@com_github_sourcegraph_conc//iter",
         "@com_github_sourcegraph_conc//pool",
         "@com_github_sourcegraph_go_diff//diff",
         "@com_github_sourcegraph_go_langserver//pkg/lsp",

cmd/frontend/graphqlbackend/cody_context.go

@@ -1,9 +1,17 @@
 package graphqlbackend
 
 import (
+	"bytes"
 	"context"
 
 	"github.com/graph-gophers/graphql-go"
+	"github.com/sourcegraph/conc/iter"
+
+	"github.com/sourcegraph/sourcegraph/cmd/frontend/internal/codycontext"
+	"github.com/sourcegraph/sourcegraph/internal/database"
+	"github.com/sourcegraph/sourcegraph/internal/gitserver"
+	"github.com/sourcegraph/sourcegraph/internal/trace"
+	"github.com/sourcegraph/sourcegraph/lib/pointers"
 )
 
 type CodyContextResolver interface {
@@ -14,6 +22,7 @@ type CodyContextResolver interface {
 	UrlMentionContext(ctx context.Context, args UrlMentionContextArgs) (UrlMentionContextResolver, error)
 	// GetCodyContext is the existing Cody Enterprise context endpoint
 	GetCodyContext(ctx context.Context, args GetContextArgs) ([]ContextResultResolver, error)
+	GetCodyContextAlternatives(ctx context.Context, args GetContextArgs) (*ContextAlternativesResolver, error)
 }
 
 type GetContextArgs struct {
@@ -33,6 +42,98 @@ type UrlMentionContextResolver interface {
 	Content() string
 }
 
+func NewContextAlternativesResolver(db database.DB, gitserverClient gitserver.Client, contextAlternatives *codycontext.GetCodyContextResult) *ContextAlternativesResolver {
+	return &ContextAlternativesResolver{
+		db:                  db,
+		gitserverClient:     gitserverClient,
+		ContextAlternatives: contextAlternatives,
+	}
+}
+
+type ContextAlternativesResolver struct {
+	db                  database.DB
+	gitserverClient     gitserver.Client
+	ContextAlternatives *codycontext.GetCodyContextResult
+}
+
+func (c *ContextAlternativesResolver) ContextLists() []*ContextListResolver {
+	var res []*ContextListResolver
+	for _, contextList := range c.ContextAlternatives.ContextLists {
+		res = append(res, &ContextListResolver{ContextList: &contextList, db: c.db, gitserverClient: c.gitserverClient})
+	}
+	return res
+}
+
+type ContextListResolver struct {
+	db              database.DB
+	gitserverClient gitserver.Client
+	ContextList     *codycontext.ContextList
+}
+
+func (r *ContextListResolver) ContextItems(ctx context.Context) (_ []ContextResultResolver, err error) {
+	tr, ctx := trace.New(ctx, "resolveChunks")
+	defer tr.EndWithErr(&err)
+
+	return iter.MapErr(r.ContextList.FileChunks, func(fileChunk *codycontext.FileChunkContext) (ContextResultResolver, error) {
+		return r.fileChunkToResolver(ctx, fileChunk)
+	})
+}
+
+func (r *ContextListResolver) Name(ctx context.Context) (string, error) {
+	return r.ContextList.Name, nil
+}
+
+// The rough size of a file chunk in runes. The value 1024 is due to historical reasons -- Cody context was once based
+// on embeddings, and we chunked files into ~1024 characters (aiming for 256 tokens, assuming each token takes 4
+// characters on average).
+//
+// Ideally, the caller would pass a token 'budget' and we'd use a tokenizer and attempt to exactly match this budget.
+const chunkSizeRunes = 1024
+
+func (r *ContextListResolver) fileChunkToResolver(ctx context.Context, chunk *codycontext.FileChunkContext) (ContextResultResolver, error) {
+	repoResolver := NewMinimalRepositoryResolver(r.db, r.gitserverClient, chunk.RepoID, chunk.RepoName)
+
+	commitResolver := NewGitCommitResolver(r.db, r.gitserverClient, repoResolver, chunk.CommitID, nil)
+	stat, err := r.gitserverClient.Stat(ctx, chunk.RepoName, chunk.CommitID, chunk.Path)
+	if err != nil {
+		return nil, err
+	}
+
+	gitTreeEntryResolver := NewGitTreeEntryResolver(r.db, r.gitserverClient, GitTreeEntryResolverOpts{
+		Commit: commitResolver,
+		Stat:   stat,
+	})
+
+	// Populate content ahead of time so we can do it concurrently
+	content, err := gitTreeEntryResolver.Content(ctx, &GitTreeContentPageArgs{
+		StartLine: pointers.Ptr(int32(chunk.StartLine)),
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	numLines := countLines(content, chunkSizeRunes)
+	endLine := chunk.StartLine + numLines - 1 // subtract 1 because endLine is inclusive
+	return NewFileChunkContextResolver(gitTreeEntryResolver, chunk.StartLine, endLine), nil
+}
+
+// countLines finds the number of lines corresponding to the number of runes. We 'round down'
+// to ensure that we don't return more characters than our budget.
+func countLines(content string, numRunes int) int {
+	if len(content) == 0 {
+		return 0
+	}
+
+	if content[len(content)-1] != '\n' {
+		content += "\n"
+	}
+
+	runes := []rune(content)
+	truncated := runes[:min(len(runes), numRunes)]
+	in := []byte(string(truncated))
+	return bytes.Count(in, []byte("\n"))
+}
+
 type ContextResultResolver interface {
 	ToFileChunkContext() (*FileChunkContextResolver, bool)
 }

cmd/frontend/graphqlbackend/cody_context.graphql

@@ -1,6 +1,6 @@
 extend type Query {
     """
-    Get pieces of context related to the query from a set of repositories.
+    Get a list of context related to the query from a set of repositories. This is the first list from getCodyContextAlternatives
     """
     getCodyContext(
         """
@@ -27,6 +27,34 @@ extend type Query {
         textResultsCount: Int!
     ): [CodyContextResult!]!
 
+    """
+    Get lists of context related to the query from a set of repositories.
+    """
+    getCodyContextAlternatives(
+        """
+        The repositories to search.
+        """
+        repos: [ID!]!
+        """
+        An optional list of file patterns used to filter the results. The
+        patterns are regex strings. For a file chunk to be returned a context
+        result, the path must match at least one of these patterns.
+        """
+        filePatterns: [String!]
+        """
+        A natural language query string.
+        """
+        query: String!
+        """
+        The number of code results to return.
+        """
+        codeResultsCount: Int!
+        """
+        The number of text results to return. Text results contain Markdown files and similar file types primarily used for writing documentation.
+        """
+        textResultsCount: Int!
+    ): CodyContextAlternativesResult!
+
     """
     EXPERIMENTAL: Detect intent for a given Cody chat query.
     """
@@ -132,6 +160,33 @@ or wiki page).
 """
 union CodyContextResult = FileChunkContext
 
+"""
+EXPERIMENTAL: The result from the Cody context alternatives API, which returns a list
+of context lists, each of which represents a distinct context list that could have been
+used. The first list is the one that should actually be used.
+"""
+type CodyContextAlternativesResult {
+    """
+    List of context lists
+    """
+    contextLists: [CodyContextList!]!
+}
+
+"""
+EXPERIMENTAL: A named list of context items.
+"""
+type CodyContextList {
+    """
+    Name of the list
+    """
+    Name: String!
+
+    """
+    List of context items
+    """
+    contextItems: [CodyContextResult!]!
+}
+
 """
 A piece of context that represents a chunk of a file in a git repository
 """

cmd/frontend/graphqlbackend/repository_reindex.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/graph-gophers/graphql-go"
 
+	"github.com/sourcegraph/sourcegraph/cmd/frontend/internal/search/idf"
 	"github.com/sourcegraph/sourcegraph/internal/auth"
 	"github.com/sourcegraph/sourcegraph/internal/search/zoekt"
 )
@@ -23,6 +24,10 @@ func (r *schemaResolver) ReindexRepository(ctx context.Context, args *struct {
 		return nil, err
 	}
 
+	if err := idf.Update(ctx, r.logger, repo.RepoName()); err != nil {
+		return nil, err
+	}
+
 	err = zoekt.Reindex(ctx, repo.RepoName(), repo.IDInt32())
 	if err != nil {
 		return nil, err

cmd/frontend/internal/codycontext/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
     visibility = ["//cmd/frontend:__subpackages__"],
     deps = [
         "//cmd/frontend/internal/cody",
+        "//cmd/frontend/internal/search/idf",
         "//internal/api",
         "//internal/cast",
         "//internal/conf",