Injected Context Architecture
## What is this epic?
This epic contains the initial iteration and implementation design of the [AI Context Management Epic](https://gitlab.com/groups/gitlab-org/-/epics/14372 "Duo: AI Context Management") for the [Injected File Context](https://gitlab.com/groups/gitlab-org/-/epics/14734 "AI Taskforce : Injected file context in Duo Chat for IDE") epic. All issues for engineering design and implementation will fall under this epic.
### What problem are we solving?
We are targeting bringing Context to Duo Chat in the Editors. You can find more information and product design in the parent epic: https://gitlab.com/groups/gitlab-org/-/epics/14734
# Injected Context Architecture
{width="370" height="239"}
Initially, we will allow users to inject their local workspace files as context.
There are multiple high-level components, each with sub-components that must be connected to bring this feature to life. Above, you can find a generic outline showing how all components interact. Here is a small breakdown of each of these components, as shown in the drawing above:
<table>
<tr>
<th>Component</th>
<th>Description</th>
</tr>
<tr>
<td>GitLab UI</td>
<td>
This will let the user use a slash command like `/include` and pull open a search menu to add their desired context. The Duo Chat UI is located in the [gitlab-ui](https://gitlab.com/gitlab-org/gitlab-ui) repository and will need to be updated to add these features.
</td>
</tr>
<tr>
<td>VS Code Extension</td>
<td>
The above UI components `GlDuoChat` are then imported into a standard Vue2 app, which is a webview in the [gitlab-vscode-extension](https://gitlab.com/gitlab-org/gitlab-vscode-extension).
The work involved in this component will involve the following:
* Integrating the changes from the `GlDuoChat` into the App.vue component
* Passing any applicable interactions from the vue app to extension runtime
* Passing the interactions form the extension runtime with the gitlab language server
</td>
</tr>
<tr>
<td>GitLab Language Server</td>
<td>
The GitLab language server will act as the primary place for the following:
* Retrieving context from each context provider
* Storing the user's context interactions (adding and deleting context items)
* Resolving any metadata for each context item
* in the case of file search, it will include a native gitignore implementation so that we can search a user's workspace files to add to Duo Chat
* Will utilize existing Open Tabs work wherever possible, eg the `DuoProjectAccessChecker`
</td>
</tr>
<tr>
<td>Duo Chat GraphQL API</td>
<td>
* We will pass the context derived from the user's actions above into the DuoChat Graphql API as `additional_context`. This will be sent whenever the user adds context to their context menu and sends a new message.
* The API will return the associated `AiContextItems` with each message that was added as context.
</td>
</tr>
<tr>
<td>AI Gateway</td>
<td>
* We must ensure AI Gateway accepts additional context as a parameter when handling the ReAct agent. Any prompt building will handle this extra context.
* Correct prompt will be built according to the type of context item passed.
</td>
</tr>
</table>
{width="374" height="207"}
# AiContextItem
This is the primary "contract" and secret sauce used _throughout_ the stack and request flow.
Here is the item:
```ts
type AIContextItem = {
id: string;
schemaVersion: `v0.0.1`
category: 'snippet' | 'file' | 'merge_request'; // larger category
type: 'open_tabs' | 'local_file_search' ; // the exact type of context
isEnabled: boolean;
disabledReasons?: string[];
metadata: object;
};
```
- `id` is going to be the unique ID to the item. We do not have to worry about id collision as there should be _no_ duplicate context item sent to the LLM. Eg; file URI, issue URL, snippet file URI + line numbers
- `schemaVersion`: this is so we can ensure backward compatibility across the clients (if context is ever enabled on the monolith side)
- `isEnabled` indication of whether the item is enabled or not. We should _never_ have `isEnabled` items sent to the backend (graphql)
- `disabledReasons` is an array of strings with reasons (in ordered precedence) for why the item is disabled. This is derived from the Context Management Policy.
- `metadata` see metadata section
### Metadata
This will be a free-form JSON object that will allow the UI/backend/consumer to alter their behavior based on the metadata of the context item.
For example, let's assume we have the following item:
```json
{
"id": "file://some/user/file.ts",
"schemaVersion": "v0.0.1",
"category": "file",
"type": "open_tabs",
"isEnabled": true,
"metadata": {
"relativeFilePath": "user/file.ts",
"project": "gitlab-lsp",
}
}
```
You can notice that metadata will be subject to change depending on the item type.
As another example, here is what an MR `AiContextItem` would look like:
```json
{
"id": "gitlab_mr_789012",
"schemaVersion": "v0.0.1",
"category": "merge_request",
"type": "merge_request",
"isEnabled": false,
"disabledReasons": ["Duo is disabled for this project"],
"metadata": {
"iid": 42,
"title": "Implement new feature X",
"description": "We need to add feature X to improve user experience",
"author": "jane.doe@example.com",
"created_at": "2024-08-25T09:15:00Z",
"labels": ["enhancement", "priority-medium"]
}
}
```
This will allow the UI to update dynamically according to the context item type. We'd want to show the `iid` in the UI in this case.
# Context Client API Design
This design outlines the various events triggered by user interactions that must be handled by the Duo Chat UI, VSCode, and Language Server.
> Note: this API design is following the conventional JSON RPC pattern of communication between VSCode and Language Server. When we migrate to a unified webview, we will need to ensure we backport this logic to the new architecture.
### DuoChat UI -\> VSCode Extension
The first communication layer we need to build is the generalized UI of the VSCode extension `App.vue`.
Here are the events and their description:
* _(Component emit -\> App.vue)_ `toggle_context_menu`: This event will be emitted on the Duo Chat component mount to the App.vue when the user requests to open the menu. The App.vue will then _request_ the enabled categories to the extension and language server. This enables the [unit primitives epic](https://gitlab.com/groups/gitlab-org/-/work_items/14979#note_2075813135 "Create a Unit Primitive for the include-file context and make it part of Duo Pro"). Eg: the UI will populate "Files, Merge Requests, Symbols, Issues" as context types in the below component state.
* _(App.vue -\> Component)_ `contextCategories`: This will be an array of provider types passed to the component as a result of the above request to VSCode and LS.
* _(Component emit -\> App.vue)_ `context_item_search_query`: This event will be emitted when the user requests a search by typing in the search input field. This will request the language server.
* _(App.vue -\> Component)_ `contextItemSearchResults`: This will be an array of search results returned to the component.
* _(Component emit -\> App.vue)_ `context_item_added`: This event will be emitted when a user adds a context item.
* _(Component emit -\> App.vue)_ `context_item_removed`: This event will be emitted when a user removes a context item.
* _(Component -\> App.vue)_ `NewPrompt`: This event will be sent when the user sends a new chat message.
### VSCode Extension -\> Language Server
The VSCode Extension will communicate with the Language Server using the following requests:
* `$gitlab/ai-context/current-context-items`: To get the current context items on component mount.
* `$gitlab/ai-context/get-provider-types`: To get the available provider types for the context menu.
* `$gitlab/ai-context/query`: To search for context items based on user input.
* `$gitlab/ai-context/add`: To add a new context item.
* `$gitlab/ai-context/remove`: To remove a context item.
* `$gitlab/ai-context/retrieve`: To retrieve context items when sending a new chat message.
### Additional Details
* The VSCode Extension mostly acts as an intermediary between the DuoChat UI (App.vue and GlDuoChat component) and the Language Server.
* Error handling will need to be implemented for context item operations (search, add, remove) to provide appropriate feedback to the user.
* AI Messages may contain `AiContextItems`, which will be displayed in the GlDuoChat component when present.
### Request Flow
```mermaid
sequenceDiagram
participant GC as GlDuoChat
participant AV as App.vue
participant VE as VSCode Extension
participant LS as Language Server
participant API as GraphQL API
Note over GC,LS: 1. On component mount
GC->>VE: getCurrentContextItems
VE->>LS: $gitlab/ai-context/current-context-items
LS-->>VE: Items (if any)
VE-->>GC: currentContextItemsResult
Note over GC,LS: 2. Toggle context menu
GC->>AV: toggle_context_menu
AV->>VE: getContextMenuTypes
VE->>LS: $gitlab/ai-context/get-provider-types
LS-->>VE: contextCategories
VE-->>GC: contextCategoriesResult
Note over GC,LS: 3. Context item search
GC->>AV: context_item_search_query
AV->>VE: contextItemSearchQuery
VE->>LS: $gitlab/ai-context/query
alt No error
LS-->>VE: AiContextItem[]
VE-->>GC: contextItemSearchResult
else Error
LS-->>VE: Error
VE-->>GC: contextItemSearchError
end
Note over GC,LS: 4. Context item added
GC->>AV: context_item_added
AV->>VE: contextItemAdded
VE->>LS: $gitlab/ai-context/add
alt No error
LS-->>VE: AiContextItem
VE-->>GC: contextItemAdded
else Error
LS-->>VE: Error
VE-->>GC: contextItemError
end
Note over GC,LS: 5. Context item removed
GC->>AV: context_item_removed
AV->>VE: contextItemRemoved
VE->>LS: $gitlab/ai-context/remove
alt No error
LS-->>VE: AiContextItem
VE-->>GC: contextItemRemoved
else Error
LS-->>VE: Error
VE-->>GC: contextItemError
end
Note over GC,API: 6. New chat message
GC->>VE: NewPrompt
VE->>LS: $gitlab/ai-context/retrieve
LS-->>VE: Context items (if any)
VE->>API: Chat message with context
Note over GC,API: 7. Receiving AI Messages
API-->>VE: AiMessages
VE-->>GC: AiMessages
Note over GC: Display AiContextItems if present
```
Describing the above:
<details>
<summary>Click to expand</summary>
1. On component mount, `GlDuoChat` sends `getCurrentContextItems` to VSCode Extension
1. VSCode Extension requests `$gitlab/ai-context/current-context-items` to Language Server
2. If there are items, VSCode Extension sends `currentContextItemsResult` to VSCode Extension
2. when `GlDuoChat` component `toggle_context_menu` emitted
1. App.vue passes `getContextMenuTypes` to VSCode Extension
2. VSCode Extension requests `$gitlab/ai-context/get-provider-types` to Language Server
3. `contextCategories` are returned to the VSCode Extension
4. VSCode extension sends `contextCategoriesResult` to `GlDuoChat` component
3. when `GlDuoChat` component `context_item_search_query` emitted
1. App.vue passes `contextItemSearchQuery` to VSCode Extension
2. VSCode Extension requests `$gitlab/ai-context/query` to Language Server
* if no error
* `AiContextItem[]` is returned back to the VSCode extension
* VSCode Extension sends `contextItemSearchResult` to the `GlDuoChat` component
* if error
* `Error` is returned back to the VSCode Extension
* VSCode Extension sends `contextItemSearchError` to the `GlDuoChat` component
4. when `GlDuoChat` component `context_item_added` emitted
1. App.vue passes `contextItemAdded` to VSCode Extension
2. VSCode Extension requests `$gitlab/ai-context/add`
* if no error
* `AiContextItem` is returned back to the VSCode extension
* VSCode Extension sends `contextItemAdded` to the `GlDuoChat` component
* if error
* `Error` is returned back to the VSCode Extension
* VSCode Extension sends `contextItemError` to the `GlDuoChat` component
5. Same for the `context_item_removed` interaction
6. when user sends a new chat message,
1. `GlDuoChat` will send `NewPrompt` to VSCode Extension
2. VSCode extension will request `$gitlab/ai-context/retrieve`
3. if there are context items returned, VSCode Extension will add those items to the chat record
4. VSCode extension send chat message with context to GraphQL API
7. when receiving `AiMessages` from GraphQL API
1. `AiMessages` are passed to `GlDuoChat` component
2. If an `AiMessage` has `AiContextItems` , they will be displayed
</details>
# GraphQL API Design
The GraphQL API will be extended to incorporate the `AiContextItem` as part of the `AiMessage` type. This addition will allow for the inclusion of context information with each AI message.
## Updated Types
### AiContextItem
Here we will define the `AiContextItem` type:
```graphql
type AiContextItem {
id: ID!
schemaVersion: String!
category: String!
type: AiContextType!
isEnabled: Boolean!
disabledReasons: [String!]
metadata: JSON!
}
```
### Updated AiMessage
We will update the `AiMessage` type to include the `AiContextItem`:
```graphql
type AiMessage {
agentVersionId: AiAgentVersionID
chunkId: Int
content: String
contentHtml: String
errors: [String!]
extras: AiMessageExtras
id: ID!
requestId: String
role: AiMessageRole!
timestamp: Time!
type: AiMessageType
contextItems: [AiContextItem!] # New field
}
```
### Updated Mutation
```
mutation {
aiAction(
input: {
chat: {
resourceId: "gid://gitlab/User/19601131",
content: "Explain this code",
additionalContext: [
{
category: SNIPPET,
id: "some_content.rb",
content: "puts 'hey'",
metadata: {}
}
]
}
}
) {
clientMutationId
errors
}
}
```
### Example Query and Response
Here's an example query to fetch AI messages with context items:
```graphql
query {
aiMessages(first: 1) {
edges {
node {
id
content
contentHtml
role
timestamp
contextItems {
id
schemaVersion
category
type
isEnabled
metadata
}
extras {
hasFeedback
sources
contextItemIds
}
}
}
}
}
```
Example response:
```json
{
"data": {
"aiMessages": {
"edges": [
{
"node": {
"id": "gid://gitlab/AiMessage/1",
"content": "Here's an explanation of the code you provided...",
"contentHtml": "<p>Here's an explanation of the code you provided...</p>",
"role": "ASSISTANT",
"timestamp": "2024-08-27T12:34:56Z",
"contextItems": [
{
"id": "file://some/user/file.ts",
"schemaVersion": "v0.0.1",
"category": "FILE",
"type": "open_tabs",
"isEnabled": true,
"metadata": {
"relativeFilePath": "user/file.ts",
"project": "gitlab-lsp"
}
}
],
"extras": {
"hasFeedback": false,
"sources": [],
"contextItemIds": ["file://some/user/file.ts"]
}
}
}
]
}
}
}
```
## Additional Considerations
1. **History and Context Persistence**: The API will maintain the history of previously submitted AI context. This can be achieved by storing the `contextItems` for each message in the conversation. When fetching a conversation history, the API will return all relevant `AiContextItem`s for each message.
5. **Context Item Versioning**: The `schemaVersion` field in `AiContextItem` allows for future extensions of the context item structure. The API should implement version checking to ensure compatibility between the client and server representations of context items.
# File Search in Language Server
### Overview
This implementation aims to provide a native, spec-compliant JavaScript implementation for `gitignore` in the GitLab language server. The goal is to improve performance when handling large workspaces with multiple repositories significantly.
##### Key Features
1. **Repository Detection**: Detects git repositories by looking for `.git/config` files.
2. **File Management**: Keeps track of files within each repository, including their ignored status.
3. **Trie-based Repository Lookup**: Uses a trie data structure for efficient repository lookups based on file paths.
4. **GitIgnore Handling**: Implements gitignore rules using the `IgnoreManager` and `IgnoreTrie` classes.
## Components
### Virtual File System Service
The Virtual File System Service utilizes `fdir` to find all files in a user's workspace and introduces `chokidar` to subscribe to workspace file changes. It employs event emitters to notify consumers of file system events.
#### Event Types
```typescript
export enum VirtualFileSystemEvents {
WorkspaceFileUpdate = 'workspaceFileUpdate',
WorkspaceFilesUpdate = 'workspaceFilesUpdate',
}
```
- `WorkspaceFileUpdate`: Emitted by `chokidar` on non-web IDE platforms when a file is changed, added, or updated.
- `WorkspaceFilesUpdate`: Emitted by `fdir` on non-web IDE platforms when the user changes their workspace configuration or when the Language Server initializes.
> Note: Web IDE support is out of scope for this iteration and will require a custom VFS implementation.
### Native GitIgnore Implementation
The native gitignore implementation is built using several classes:
1. **IgnoreTrie**: A trie-based data structure for efficient storage and lookup of gitignore patterns.
- Supports adding patterns at specific paths.
- Efficiently checks if a given path is ignored.
2. **IgnoreManager**: Manages the loading and application of gitignore rules for a repository.
- Loads gitignore files from the repository.
- Uses IgnoreTrie to store and check ignore patterns.
3. **Repository**: Integrates the IgnoreManager to handle file ignoring within a repository.
- Manages a collection of files and their ignored status.
- Provides methods to add, remove, and query files in the repository.
### Repository Service
The Repository Service is responsible for managing git repositories within the workspace. It listens to events from the Virtual File System Service and maintains a structure of repositories and their files.
##### Main Classes
- `Repository`: Represents a git repository, managing its files and ignore rules.
- `RepositoryTrie`: A trie data structure for efficient repository lookups.
- `DefaultRepositoryService`: The main service class that coordinates repository management.
### Benefits of this Architecture
1. **Scalability**: Significant performance boost compared to `isomorphic-git`'s `getStatusMatrix` method, which took over 13 seconds for the same task on the GitLab Monolith
2. **Real-time Updates**: Uses `chokidar` to watch for file system changes and update the internal state accordingly.
3. **Nested Repository Support**: Handles workspaces with multiple nested git repositories.
4. **Extensibility**: The `repository_service` is designed to be general-purpose, allowing for future integration with other git-related tasks.
5. **Security**: By not calling git directly, it reduces the risk of path execution security issues.
### Implementation Details
- The system uses URI objects for consistent file path handling across different platforms.
- A trie-based structure (`RepositoryTrie`) will be used for efficient repository lookups based on file paths.
- The `IgnoreTrie` class provides an efficient way to store and check gitignore patterns, improving performance for large repositories with complex ignore rules.
- The `DefaultRepositoryService` class acts as the main orchestrator, handling workspace updates, file changes, and providing an interface for querying repository information.
### Future Considerations
1. Web IDE Support: Implement a custom VFS for web-based environments.
2. Migrate `DuoProjectAccessCache` to this new architecture for Open Tabs functionality.
###
epic