sourcegraph/dev
Rok Novosel f77f0272cf
embeddings: searcher and indexer (#48017)
# High-level architecture overview
<img width="2231" alt="Screenshot 2023-02-24 at 15 13 59"
src="https://user-images.githubusercontent.com/6417322/221200130-53c1ff25-4c47-4532-885f-5c4f9dadb05e.png">


# Embeddings

Really quickly: embeddings are a semantic representation of text.
Embeddings are usually floating-point vectors with 256+ elements. The
neat thing about embeddings is that they allow us to search over textual
information using a semantic correlation between the query and the text,
not just syntactic (matching keywords).

In this PR, we implemented an embedding service that will allow us to do
semantic code search over repositories in Sourcegraph. So, for example,
you'll be able to ask, "how do access tokens work in Sourcegraph", and
it will give you a list of the closest matching code files.

Additionally, we build a context detection service powered by
embeddings. In chat applications, it is important to know whether the
user's message requires additional context. We have to differentiate
between two cases: the user asks a general question about the codebase,
or the user references something in the existing conversation. In the
latter case, including the context would ruin the flow of the
conversation, and the chatbot would most likely return a confusing
answer. We determine whether a query _does not_ require additional
context using two approaches:

1. We check if the query contains well-known phrases that would indicate
the user is referencing the existing conversation (e.g., translate
previous, change that)
1. We have a static dataset of messages that require context and a
dataset of messages that do not. We embed both datasets, and then, using
embedding similarity, we can check which set is more similar to the
query.

## GraphQL API

We add four new resolvers to the GraphQL API:

```graphql
extend type Query {
  embeddingsSearch(repo: ID!, query: String!, codeResultsCount: Int!, textResultsCount: Int!): EmbeddingsSearchResults!
  isContextRequiredForQuery(query: String!): Boolean!
}
extend type Mutation {
  scheduleRepositoriesForEmbedding(repoNames: [String!]!): EmptyResponse!
  scheduleContextDetectionForEmbedding: EmptyResponse!
}
```

- `embeddingsSearch` performs embeddings search over the repo embeddings
and returns the specified number of results
- `isContextRequiredForQuery` determines whether the given query
requires additional context
- `scheduleRepositoriesForEmbedding` schedules a repo embedding
background job
- `scheduleContextDetectionForEmbedding` schedules a context detection
embedding background job that embeds a static dataset of messages.

## Repo embedding background job

Embedding a repository is implemented as a background job. The
background job handler receives the repository and the revision, which
should be embedded. Handler then gathers a list of files from the
gitserver and excludes files >1MB in size. The list of files is split
into code and text files (.md, .txt), and we build a separate embedding
index for both. We split them because in a combined index, the text
files always tended to feature as top results and didn't leave any room
for code files. Once we have the list of files, the procedure is as
follows:

- For each file
  - Get file contents from gitserver
- Check if the file is embeddable (is not autogenerated, is large
enough, does not have long lines)
  - Split the file into embeddable chunks
- Embed the file chunks using an external embedding service (defined in
site config)
  - Add embedded file chunks and metadata to the index
- Metadata contains the file name, the start line, and the end line of
the chunk
- Once all files are processed, the index is marshaled into JSON and
stored in Cloud storage (GCS, S3)

### Site config changes

As mentioned, we use a configurable external embedding API that does the
actual text -> vector embedding part. Ideally, this allows us to swap
embedding providers in the future.

```json
"embeddings": {
  "description": "Configuration for embeddings service.",
  "type": "object",
  "required": ["enabled", "dimensions", "model", "accessToken", "url"],
  "properties": {
    "enabled": {
      "description": "Toggles whether embedding service is enabled.",
      "type": "boolean",
      "default": false
    },
    "dimensions": {
      "description": "The dimensionality of the embedding vectors.",
      "type": "integer",
      "minimum": 0
    },
    "model": {
      "description": "The model used for embedding.",
      "type": "string"
    },
    "accessToken": {
      "description": "The access token used to authenticate with the external embedding API service.",
      "type": "string"
    },
    "url": {
      "description": "The url to the external embedding API service.",
      "type": "string",
      "format": "uri"
    }
  }
}
```

## Repo embeddings search

The repo embeddings search is implemented in its own service. When a
user queries a repo using embeddings search, the following happens:

- Download the repo embedding index from blob storage and cache it in
memory
  - We cache up to 5 embedding indexes in memory
- Embed the query and use the embedded query vector to find similar code
and text file metadata in the embedding index
- Query gitserver for the actual file contents
- Return the results

## Interesting files

- [Similarity
search](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-102cc83520004eb0e2795e49bc435c5142ca555189b1db3a52bbf1ffb82fa3c6)
- [Repo embedding job
handler](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-c345f373f426398beb4b9cd5852ba862a2718687882db2a8b2d9c7fbb5f1dc52)
- [External embedding api
client](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-ad1e7956f518e4bcaee17dd9e7ac04a5e090c00d970fcd273919e887e1d2cf8f)
- [Embedding a
repo](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-1f35118727128095b7816791b6f0a2e0e060cddee43d25102859b8159465585c)
- [Embeddings searcher
service](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-5b20f3e7ef87041daeeaef98b58ebf7388519cedcdfc359dc5e6d4e0b021472e)
- [Embeddings
search](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-79f95b9cc3f1ef39c1a0b88015bd9cd6c19c30a8d4c147409f1b8e8cd9462ea1)
- [Repo embedding index cache
management](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-8a41f7dec31054889dbf86e97c52223d5636b4d408c6b375bcfc09160a8b70f8)
- [GraphQL
resolvers](https://github.com/sourcegraph/sourcegraph/pull/48017/files#diff-9b30a0b5efcb63e2f4611b99ab137fbe09629a769a4f30d10a1b2da41a01d21f)


## Test plan

- Start by filling out the `embeddings` object in the site config (let
me know if you need an API key)
- Start the embeddings service using `sg start embeddings`
- Go to the `/api/console` page and schedule a repo embedding job and a
context detection embedding job:

```graphql
mutation {
  scheduleRepositoriesForEmbedding(repoNames: ["github.com/sourcegraph/handbook"]) {
    __typename
  }
  scheduleContextDetectionForEmbedding {
    __typename
  }
}
```

- Once both are finished, you should be able to query the repo embedding
index, and determine whether context is need for a given query:

```graphql
query {
  isContextRequiredForQuery(query: "how do access tokens work")
  embeddingsSearch(
    repo: "UmVwb3NpdG9yeToy", # github.com/sourcegraph/handbook GQL ID
    query: "how do access tokens work", 
    codeResultsCount: 5,
    textResultsCount: 5) {
    codeResults {
      fileName
      content
    }
    textResults {
      fileName
      content
    }
  }
}
```
2023-03-01 10:50:12 +01:00
..
adr-docs bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
auth-provider Remove left-over mentions of dev/start.sh (#25745) 2021-10-07 16:38:00 +02:00
authtest Fix integration test for list users (#48276) 2023-02-27 15:11:13 +00:00
bkstats bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
build-tracker Use new base image with curl from edge/main (#48144) 2023-02-23 19:40:07 +00:00
buildchecker Migrate to autogold/v2 (needed by Bazel) (#47891) 2023-02-21 10:37:13 +01:00
check embeddings: searcher and indexer (#48017) 2023-03-01 10:50:12 +01:00
ci packages: stage 2 of packages tables migration (#48007) 2023-02-23 15:21:15 +00:00
codeintel-qa codeintel: Remove oobmigration step from codeintel-qa (#47714) 2023-02-17 13:45:00 -06:00
corrupt-archives bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
db dev/deb: Remove scripts replaced by sg (#28767) 2021-12-08 23:22:00 +00:00
depgraph Housekeeping: Add package name aliases to avoid collisions with variables (#47180) 2023-01-31 16:28:57 +01:00
dx bazel: build //enterprise (#47327) 2023-02-07 15:30:46 +01:00
gqltest gqltest: output unexpected repos in context test (#48249) 2023-02-28 08:33:20 +02:00
grafana dev/sg: introduce 'sg ci logs' with Loki support (#25835) 2021-10-12 11:14:50 -04:00
internal/cmd/auth-proxy-http-header bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
nix nix: add p4-fusion building (#44588) 2022-11-23 16:06:28 +00:00
okay Housekeeping: Remove redundant stuff from Go code (#47104) 2023-01-30 20:39:40 +00:00
perforce Migrate to autogold/v2 (needed by Bazel) (#47891) 2023-02-21 10:37:13 +01:00
perforce-testing-helpers prettier (#40267) 2022-08-11 10:26:35 -06:00
phabricator web: migrate from yarn to pnpm (#46143) 2023-01-11 19:50:09 -08:00
pr-auditor bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
prometheus otel: add collector dashboard (#45009) 2022-12-19 13:18:51 +01:00
release release-tool: auto save release config when prompting for input (#48135) 2023-02-28 12:50:56 -07:00
scaletesting Backend: replace lib/group with sourcegraph/conc (#48162) 2023-02-24 14:35:54 -07:00
sg sg: improve gen failure output and ignore vendor dir (#48261) 2023-02-27 10:36:35 +00:00
src-expose Use new base image with curl from edge/main (#48144) 2023-02-23 19:40:07 +00:00
team bazel: introduce build files for Go (#46770) 2023-01-23 14:00:01 +01:00
tilt Remove left-over mentions of dev/start.sh (#25745) 2021-10-07 16:38:00 +02:00
zoekt zoekt: set -indexserver_proxy for webserver (#44995) 2022-12-02 08:38:44 +01:00
.gitignore dev: remove nginx (#13299) 2021-02-18 14:38:58 +00:00
add_https_domain_to_hosts.sh dev: add shfmt for shell script consistency (#9900) 2020-04-15 12:44:36 -07:00
babel.bzl bazel: add bazel build,tests for client/* (#46193) 2023-02-28 20:46:03 -08:00
BUILD.bazel bazel: add bazel build,tests for client/* (#46193) 2023-02-28 20:46:03 -08:00
caddy.sh Bump caddy version (#26329) 2021-10-19 22:06:22 +00:00
Caddyfile dev/Caddyfile: use 127.0.0.1 instead of localhost (#38955) 2022-07-18 13:19:09 -07:00
CLA.txt
codecov.yml vscode: ignore vscode in Codecov (#32676) 2022-03-16 11:28:28 -04:00
codeinsights-db.sh insights: update references to TimescaleDB (#32948) 2022-03-29 11:04:47 +01:00
comby-install-or-upgrade.sh search: bump comby 1.8.1 (#37804) 2022-06-28 11:39:49 -07:00
ctags-install.sh dev: add installer script for building a local universal-ctag (#45198) 2022-12-06 18:32:21 +01:00
defs.bzl bazel: add bazel build,tests for client/* (#46193) 2023-02-28 20:46:03 -08:00
dev-sourcegraph-server.sh ci: add shellcheck linter for shell scripts (#9903) 2020-04-21 10:03:17 -07:00
docsite.sh Upgrade docsite version to 1.9.1 in dev/docsite.sh (#47663) 2023-02-21 09:50:58 +01:00
drop-entire-local-database-and-redis.sh all: /bin/bash -> /usr/bin/env bash (#23673) 2021-08-06 12:02:43 +02:00
foreach-ts-project.sh [SG-46115] - Merge the search-ui package into the branded package (#46197) 2023-01-11 18:27:19 +01:00
git-stats search: create and document git-stats script (#32663) 2022-03-16 13:41:17 +02:00
global-settings.json use new non-extensions-based panel, remove old HierarchicalLocationsView (#46052) 2023-02-10 10:34:49 -08:00
go-mod-update.sh Upgrade aws-sdk-go-v2 (#19155) 2021-04-14 15:06:15 +02:00
golangci-lint.sh dev: download golangci-lint instead of building it (#45259) 2022-12-06 17:52:58 +01:00
licenses.sh rework plugin structure and implement frontside blogpost (#46883) 2023-02-15 11:49:51 +02:00
mocha.bzl bazel: add bazel build,tests for client/* (#46193) 2023-02-28 20:46:03 -08:00
postgres_exporter.sh migrations: Update postgres exporter queries for dirty database (#30774) 2022-02-10 17:40:09 -06:00
prune-pick.sh ci: add shellcheck linter for shell scripts (#9903) 2020-04-21 10:03:17 -07:00
redis-postgres.yml Update Postgres to 12.7 (#31933) 2022-03-03 14:39:52 -06:00
redis.conf Local dev with docker-compose (#23537) 2021-08-06 22:11:40 +00:00
run-server-image.sh gRPC: run tests with and without gRPC (#47770) 2023-02-17 12:38:41 -07:00
sass.bzl build: move build-config specific dependencies to build-config/package.json (#48072) 2023-02-22 19:25:07 -08:00
site-config.json Always disable user external service mode in UI (#44721) 2022-11-23 14:49:03 +01:00
src-prof-services.json Push executor metrics (#36969) 2022-08-03 12:08:04 +02:00
src-search-meta.sh GraphQL: remove deprecated resultCount field (#31573) 2022-02-21 12:03:36 -07:00
tools.go unrevert "all: rename go module google/zoekt to sourcegraph/zoekt" (#40423) 2022-08-17 15:30:49 -07:00
universal-ctags-dev dev: add installer script for building a local universal-ctag (#45198) 2022-12-06 18:32:21 +01:00
webpack.bzl bazel: add bazel build,tests for client/* (#46193) 2023-02-28 20:46:03 -08:00