**DRIs: @mkaeppler (Cloud Connector), @shinya.maeda (AI Framework)** ## Problem Duo Chat currently relies on embeddings for parts of its functionality (currently: GitLab Docs). The way this process works for SaaS is as follows: 1. Documentation markdown files are chunked and stored in a Postgres database. This happens every day in a Sidekiq job, which is accounting for the fact that for SaaS, docs are constantly in flux. 1. When a user asks a question, it is sent to the AI model to turned it into an embedding 1. The embeddings storage is searched for matching embeddings using the question embedding and a vector proximity search, and the associated text is returned 1. The content retrieved this way is sent as context along with the prompt to the model to produce the final answer This can be visualized with this sequence diagram: ```mermaid sequenceDiagram autonumber participant U as User participant GL as GitLab participant DB as Embeddings Database participant M as AI model Note over GL,M: Cloud-managed loop Sidekiq: Update docs embeddings GL->>GL: Parse and chunk docs markdown GL->>M: Fetch embeddings M-->>GL: Docs embeddings GL->>DB: Store embeddings end U->>GL: Send question GL->>M: Encode question M-->>GL: Question embedding GL->>DB: Proximity search with question embedding DB-->>GL: N closest text matches GL->>M: Send prompt with text matches M-->>GL: Final answer GL-->>U: Final answer ``` In order to bring this to self-managed/Dedicated, we need to address the following questions: 1. **Where are embeddings stored?** We could use storage that is local to the GitLab instance, or host this database on behalf of customers. 1. **How are embeddings stored?** We currently use a dedicated Postgres database to store embeddings. In order to support vector based search queries, it requires the `pgvector` extension to be installed, which is not the case for our default PG setup for SM. - We decided to use the same storage combo for SM. We verified that this extension is compliant with GL licensing, we can ship it to customers, and is supported by all cloud providers we support as per [this comment](https://gitlab.com/gitlab-org/gitlab/-/issues/404396#note_1566856845). We furthermore discarded the option of using Cloud SQL as per [this comment](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423210). 1. **How are embeddings populated?** Regardless of whether we use local or remote storage, an open question is how this database would be populated. This step requires parsing GitLab docs markdown and making calls into the AI model to turn them into embeddings. This work is subject to AI vendor quotas. 1. **AI model support:** We currently support OpenAI embeddings with VertexAI support being added but still experimental. An open question was whether we need to support both for SM too. - We currently expect that we will complete the work to [move away from OpenAI embeddings](https://gitlab.com/gitlab-org/gitlab/-/issues/420939) before finishing the work described here. This work is handled by gitlab~29783025. ## Solution exploration ### Descoped: dynamic and private data Generally, we can think of data broadly in the following dimensions: | Nature | Availability | | ------- | ------------ | | static | public | | dynamic | private | Embeddings won't just be used for GitLab docs, which are **static/public** in nature. They will also be [needed for MRs, issues, source code](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423212) etc. i.e. data that is different for each customer and may be private in nature i.e. not allowed to leave their instances. Related discussion [in this thread](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423244). ### Constraints We need to consider the following constraints for working out a solution: - **Storage size.** If we decide to ship some sort of embeddings artifact to customers (a pre-seeded DB or an intermediate format used to import embeddings), or directly produce embeddings in the customer instance, we are constrained by database growth as [mentioned here](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423412). - **Implication:** Any solution that involves customers using/importing such an artifact from us is only an option if it does not grow their local database by more than X%. - **AI model quotas.** Any solution will require us to call into the AI model to retrieve embeddings. For the existing SaaS solution we are currently constrained to 600 RPM as [mentioned here](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423382). Our goal should be to not materially add to this request volume, or ask to get this raised. - **Implication:** Any solution that sees the customer call into the AI model to produce the embeddings themselves is infeasible since it would imply thousands of instances redundantly producing essentially the same data for their individual use, which would quickly burn through our 3P vendor quota. ### Solution dimensions This is a multi-dimensional problem since there are various ways to handle each aspect outlined in the problem statement. This section briefly summarizes these dimensions. #### Local vs remote storage Embeddings storage could live either on premises or be hosted by us. This gives rise to the following options: 1. **Embeddings database is hosted by customer** 1. **Create embeddings build artifact during release, ship it to customers.** We could either pre-seed a database from the documentation text as it was current at the time of release, or create some other intermediate representation (e.g. JSON dump) that we then bundle and ship with a milestone release. Ideally this only happens once at the time we promote a release since it requires talking to the AI model to obtain embeddings. We identified [challenges with this](https://gitlab.com/groups/gitlab-org/-/epics/11600#note_1583423289) since our release pipeline produces immutable packages, so at the time this happens we cannot include anything else in the package anymore. 2. **Create embeddings build artifact during release, make it available for import.** Alternatively, we could create this artifact, but instead of bundling it with the release, make it available for download somewhere so that SM instance can import it, either during the upgrade process or in response to an application trigger (e.g. enabling the Chat feature.) 1. **Embeddings database is hosted by us** 1. **Serve embeddings from SaaS database.** We already import embeddings on a nightly basis for SaaS. We could make this data available to SM too. The main problems to solve here would be that we would have to start versioning this data since SM instances need a specific fixed-time view on this data, and we would have to make it available through an API so the application can request it. 2. **Serve embeddings from dedicated database.** Alternatively, we could host a dedicated Postgres embeddings database for Cloud Connector customers. This database could be populated with similar mechanisms as outlined under `Embeddings are shipped to customer`, or be produced as a snapshot/dump from the SaaS DB at the time of release. ### Examples: Prospective solutions (not complete) Note: this is not a one-dimensional problem so there are valid permutations of some of the solutions outlined here and they are not all listed in detail. #### Approach 1: Local database at customer site + AI gateway embeddings endpoint In this approach we would merely push down 3P model access into the AI gateway but retain the overall "protocol" between the GitLab application and the model, i.e. the main logic remains in gitlab-rails and leaves the AI gateway be a simple proxy. This necessitates that documentation embeddings are made available on premises since the question vector is an input to the text search. It is still unclear how that would work, e.g. by making it available as a download: ```mermaid sequenceDiagram autonumber participant U as User participant GL as GitLab participant DB as Embeddings Database participant AI as AI gateway participant M as AI model Note over U,DB: Self-managed Note over AI,M: Cloud-managed U->>GL: Send question GL->>AI: Encode question AI->>M: Request question embedding M-->>AI: Question embedding AI-->>GL: Question embedding GL->>DB: Proximity search with question embedding DB-->>GL: N closest text matches GL->>AI: Send prompt with text matches AI->>M: Send prompt with text matches M-->>AI: Final answer AI-->>GL: Final answer GL-->>U: Final answer ``` **Pros:** - Low implementation complexity: Close to what we currently do for SaaS, maintains logic flow/protocol between gitlab-rails and the AI gateway - Low operational overhead for GitLab (the company): we neither host nor manage the customer's embeddings. - Works with well understood solutions like postgres + pgvector - Would work with private data for other AI use cases that require embeddings since the source data would never leave the SM instance **Cons:** - Increased release complexity: Requires us to provide some kind of "delivery mechanism" for embeddings data that is part of our monthly release process since each time we tag a new milestone release, we need to make a snapshot of the docs embeddings data available to SM instances. - Difficult to make changes: requires customers to act on fixes/updates to docs embeddings, as opposed to a GL-hosted solution where we could update embeddings without customers ever knowing. Adding a new vendor requires a new milestone release and a re-import of these data at the customer site. #### Approach 1.1: Cloud database + AI gateway embeddings/search endpoint This is a variation of `Approach 1` in that it maintains the overall communication flow between GitLab Rails and the model (with the AI gateway acting as a model proxy in between), but the embeddings database is not co-located with gitlab-rails and is instead hosted in the cloud: ```mermaid sequenceDiagram autonumber participant U as User participant GL as GitLab participant AI as AI gateway participant DB as Embeddings Database participant M as AI model Note over U,GL: Self-managed Note over AI,M: Cloud-managed U->>GL: Send question GL->>AI: Encode question AI->>M: Request question embedding M-->>AI: Question embedding AI-->>GL: Question embedding GL->>AI: Fetch N closest matches for question embedding AI->>DB: Proximity search with question embedding DB-->>AI: N closest text matches AI-->>GL: N closest text matches GL->>AI: Send prompt with text matches AI->>M: Send prompt with text matches M-->>AI: Final answer AI-->>GL: Final answer GL-->>U: Final answer ``` **Pros:** - Low process overhead: since GitLab hosts embeddings on behalf of customers, we do not need to cut embeddings release artifacts or otherwise ship embeddings as part of a release. - Easy to update: since we host embeddings, we can update them whenever, and as long as gitlab-rails communicates its version to the AI gateway, they will receive the most up to date embeddings. - Fewer moving parts for self-managed customers and no need to use pgvector or other extensions so the application can perform vector searches. **Cons:** - Incompatible with current implementation for SaaS: we either need to: 1. Use a code switch to obtain embeddings through AI gateway only for SM (but that violates the AI gateway architecture guidelines) 2. Or, we also move our embeddings for SaaS to this new database (meaning some migration work is necessary) 3. Or, we make the AI gateway talk to the existing SaaS DB (which would need to be reworked to store docs by GL version instead of rolling them over daily, plus this would introduce a circular dependency between GitLab-SaaS and the AI gateway) - Slightly complicates implementation complexity: we need to create an additional AI gateway endpoint that swaps a docs embedding for its neighest neighbors (e.g. `/embeddings/closest_set`), and gitlab-rails database queries need to be rewritten to make this API call instead. - Operational overhead for GitLab: since we run and scale this database, we shift the operational burden to our side (though arguably, this is what Cloud Connector is all about) - Does not work with private data since that cannot leave the SM instance (not relevant for docs, but implies that this solution would not be reusable for other AI use cases that require embeddings) #### Discarded approaches <details> <summary>Approach 2: GitLab-managed database + rich AI gateway chat API</summary> **Note:** we discarded this approach at least for now since it comes with higher implementation complexity, primarily because the Ruby prompt engine would have to move to the AI gateway. In this approach, we push more functionality into the GitLab infrastructure by simplifying the protocol between GitLab and the AI gateway. Here, the GitLab application only sends the original question to the AI gateway, the AI gateway then executes the internal protocol, including querying an embeddings database. It is also TBD yet how this database would be maintained for self-managed since it would require timestamped/snapshotted documentation by version. It also means it would change the AI gateway from being a stateless service to a stateful one because it now talks to connected storage: ```mermaid sequenceDiagram autonumber participant U as User participant GL as GitLab participant AI as AI gateway participant DB as Embeddings Database participant M as AI model Note over U,GL: Self-managed Note over AI,M: Cloud-managed U->>GL: Send question GL->>AI: Send question AI->>M: Request question embedding M-->>AI: Question embedding AI->>DB: Proximity search with question embedding DB-->>AI: N closest text matches AI->>M: Send prompt with text matches M-->>AI: Final answer AI-->>GL: Final answer GL-->>U: Final answer ``` </details>