Commit ab1cbd33 authored by Pam Artiaga's avatar Pam Artiaga 2️⃣ Committed by Arturo Herrero
Browse files

Add design documents for Active Context and Code Embeddings

parent b8d4bff9
Loading
Loading
Loading
Loading
+2 −2
Original line number Diff line number Diff line
@@ -4,8 +4,8 @@ status: ongoing
creation-date: "2024-09-11"
authors: [ "@dgruzd" ]
coaches: [ "@shekharpatnaik" ]
dris: []
owning-stage: "~devops::foundations"
dris: [ "@wortschi" ]
owning-stage: "~devops::ai platform"
participating-stages: []
toc_hide: true
---
+139 −0
Original line number Diff line number Diff line
---
title: "Active Context Tasks"
description: "Design document for the Active Context Tasks framework"
status: implemented
creation-date: "2026-06-23"
authors: [ "@partiaga" ]
coaches: [ ]
dris: [ "@wortschi" ]
owning-stage: "~devops::ai platform"
participating-stages: []
toc_hide: true
---

## Overview

`ActiveContext::Task` is a framework for managing long-running, asynchronous operations within the ActiveContext system. It provides a structured way to define, execute, and track complex workflows that may involve multiple sequential or dependent steps.

This is designed to handle operations like embedding model switching, embeddings field creation and backfill, and metadata synchronization that require careful orchestration and error handling.

## Components

### Task Base Class (`ActiveContext::Task[1.0]`)

The abstract base class for all task implementations and defines the interface that all subclasses must follow.

The concrete class implementations define the actual logic of the task.

**Batched Tasks:**

Tasks can be marked as `batched!` to indicate they perform work in batches and may need multiple executions to complete. For batched tasks, `completed?` must be implemented to determine when all work is done.

**Example Task class implementations**

- `Ai::ActiveContext::Tasks::BackfillEmbeddings`: Generates embeddings for existing documents using a new embedding model.
- `Ai::ActiveContext::Tasks::SyncFeatureSettings`: Synchronizes feature settings with the embedding model configuration.

### Task Model (`Ai::ActiveContext::Task`)

The ActiveRecord model that persists task information.

**Database Schema:**

- `connection_id` - Foreign key to the vector store connection
- `depends_on_id` - Foreign key to the preceding task (for [task chains](#task-chains))
- `name` - The task class name
- `status` - Current execution status (`pending`, `in_progress`, `completed`, `failed`)
- `params` - JSON parameters passed to the task
- `retries_left` - Number of retry attempts remaining (default: 3)
- `started_at` - When execution began
- `completed_at` - When execution finished
- `error_message` - Error details if the task failed

### Task Dictionary (`ActiveContext::Task::Dictionary`)

The registry of defined task classes. This provides a singleton instance for global access.

```ruby
# Find a task class by name
task_class = ActiveContext::Task::Dictionary.instance.find_by_name('Ai::ActiveContext::Tasks::BackfillEmbeddings')

# Or use the shorthand
task_class = ActiveContext::Task::Dictionary.find_by_name('Ai::ActiveContext::Tasks::BackfillEmbeddings')

# Returns an array of task class objects that have been loaded
ActiveContext::Task::Dictionary.instance.tasks
```

### Task Service (`Ai::ActiveContext::TaskService`)

Creates and manages task chains.

- `create_task(task_class, params: {}, depends_on: nil)` - Create a single task
- `create_chain(*tasks_with_params)` - Create a sequence of dependent tasks

**Example Usage:**

```ruby
service = Ai::ActiveContext::TaskService.new
service.create_chain(
  [Ai::ActiveContext::Tasks::AddEmbeddingsField, { collection: 'code', field: 'embeddings_v2', dimensions: 768 }],
  [Ai::ActiveContext::Tasks::BackfillEmbeddings, { collection: 'code', field: 'embeddings_v2' }],
  [Ai::ActiveContext::Tasks::UpdateCollectionMetadata, { collection: 'code', metadata: {...} }]
)
```

### Task Worker (`Ai::ActiveContext::TaskWorker`)

The Sidekiq worker responsible for executing created tasks.

**Execution Flow:**

1. Worker finds the next processable task record
1. Instantiates the correct task object from the task record `name`
1. Marks task record as `in_progress`
1. Calls `execute!` on the task object
1. On success: marks task record as `completed`
1. On failure: see "Error handling" details below
1. If task is batched and not completed: re-enqueues worker
1. If no more processable tasks: worker exits

**Error Handling:**

- Catches exceptions during task execution
- For each failed execution, the number of retries is decreased and:
  - Task stays `in_progress` with one fewer retry; the cron worker picks it up on its next run
  - Task is marked as `failed` and cascades to dependents (if no retries left)

**Re-enqueueing Logic:**

- For batched tasks that aren't completed, worker re-enqueues to continue processing
- After a successful non-batched task, the worker exits; the cron-scheduled invocation picks up the next one

## Task Chains

Tasks can depend on other tasks through the `depends_on` relationship. A task only becomes processable when:

- Its status is `pending` or `in_progress`, AND
- Either it has no dependency, OR its dependency has a `completed` status

This creates a directed acyclic graph (DAG) of task execution.

## Task Execution Flow Summary

1. **Task Creation** - A service creates one or more tasks by using `TaskService`
1. **Dependency Resolution** - Tasks are chained through `depends_on` relationships
1. **Processable Selection** - The `Ai::ActiveContext::Task.processable` scope finds tasks ready to execute
1. **Worker Polling** - `TaskWorker` finds the next processable task through `Ai::ActiveContext::Task.current`
1. **Execution** - `TaskWorker` invokes the task's `execute!` method
1. **Status Transition** - Task status moves from pending → `in_progress``completed`/`failed`
1. **Dependent Execution** - Once a task completes, its dependents become processable
1. **Worker Re-enqueueing** - For batched tasks, the worker re-enqueues itself; otherwise the cron-scheduled invocation picks up the next task
1. **Error Handling** - Failed tasks cascade failures to all dependents

## Planned Future Enhancements

- **Task Prioritization** - Allow high-priority tasks to execute before others
- **Parallel Execution** - Execute independent tasks concurrently
- **Task Cancellation** - Allow canceling pending or in-progress tasks
- **Conditional Tasks** - Execute tasks based on conditions or previous results
+30 −0
Original line number Diff line number Diff line
---
title: "AI Context Abstraction Layer ADR-001: Gem-Based Architecture"
description: "Decision record for implementing the AI Context Abstraction Layer as the ActiveContext Ruby gem."
toc_hide: true
---

## Status

Accepted

## Context

The AI Context Abstraction Layer is a critical component that provides a unified interface for semantic search across different vector stores (Elasticsearch, OpenSearch, PostgreSQL with pgvector). This layer needs to be:

1. Reusable across multiple GitLab services
2. Maintainable and testable in isolation
3. Clearly separated from GitLab-specific implementations
4. Extensible for future vector store backends

## Decision

We have decided to:

1. Implement the ActiveContext Abstraction Layer as a Ruby gem located in `gems/gitlab-active-context` within the GitLab Rails repository
2. Name the gem `ActiveContext`, thus all gem classes are namespaced under `ActiveContext::`
3. GitLab-specific implementations should be namespaced under `Ai::ActiveContext` (for example, `Ai::ActiveContext::Collections::Code`, `Ai::ActiveContext::References::Code`)

## References

- [ActiveContext Gem](https://gitlab.com/gitlab-org/gitlab/-/tree/master/gems/gitlab-active-context)
+41 −0
Original line number Diff line number Diff line
---
title: "AI Context Abstraction Layer ADR-002: ActiveContext Task framework"
description: "Decision record for the introduction of a framework for complex long-running operations"
toc_hide: true
---

## Status

Accepted

## Context

The ActiveContext framework needs to support long-running, asynchronous operations that may span multiple steps and require careful orchestration.

For example, switching from one embedding model to another requires:

1. Creating a new vector field
1. Backfilling embeddings for all documents
1. Updating collection metadata
1. Syncing feature settings
1. Cleaning up old fields

### Challenges with Existing Approaches

**Simple Sidekiq Workers:**

- No built-in support for task dependencies
- Difficult to handle cascading failures
- No visibility into multi-step workflows
- Hard to retry individual steps in a chain

**Orchestration Services:**

- Complex state management
- Difficult to recover from crashes
- No persistence of workflow state
- Hard to debug long-running operations

## Decision

We have decided to introduce the [**ActiveContext Task Framework**](../active_context_tasks.md) for managing long-running, asynchronous operations. It provides a structured way to define, execute, and track complex workflows that may involve multiple sequential or dependent steps.
+45 −0
Original line number Diff line number Diff line
---
title: "AI Context Abstraction Layer ADR-003: Embedding model redesign"
description: "Decision record explaining a new embedding model design"
toc_hide: true
---

## Status

Accepted

## Context

The original ActiveContext design hard-coded embedding models as embedding version hashes:

```ruby
class Ai::ActiveContext::Collections::Code # class implementing an ActiveContext Collection
  embeddings_v1: { model: 'text_embedding_005_vertex', dimensions: 768 }
  embeddings_v2: { model: 'text_embedding_004_vertex', dimensions: 512 }
end
```

### Limitations of the previous design

- Model selection by the user cannot be supported, making it a blocker for Self-Managed instances with Duo Self-hosted, where the user must select their own models
- Adding new models requires updating the hard-coded version hashes and doing a manual backfill process, creating maintenance burden

### Existing Model Configuration feature

There is already an existing [Model Configuration admin page](https://docs.gitlab.com/administration/gitlab_duo/model_selection/) for other AI features. However, this cannot support the backfill process needed when switching embedding models.

## Decision

We have decided to redesign the embedding models to support flexible model selection with one-click model switching that triggers a backfill process in the background.

The main concepts of the redesign are:

- **Model Metadata** - embedding model configuration is persisted in `Ai::ActiveContext::Collection` record `metadata` rather than hard-coded
- **3 model configurations per Collection**
  - `current_indexing_embedding_model`: The model currently used for indexing content
  - `next_indexing_embedding_model`: The model queued to replace the current model (used when switching to a new model)
  - `search_embedding_model`: The model used for query embeddings during search. This matches the `current_indexing_embedding_model`
- **Asynchronous switching process** - Model switching triggers an asynchronous process that leverages the [Active Context Tasks framework](../active_context_tasks.md)
- **Indirect syncing with AI Feature Settings** - A Collection's `current_indexing_embedding_model` will be synced with the `Ai::FeatureSetting` records used by other AI features

For further details, please see the [Embedding Models design document](../embedding_models.md).
Loading