kiln-ai 0.20.1__py3-none-any.whl → 0.22.0__py3-none-any.whl

kiln_ai/adapters/rag/rag_runners.py

@@ -0,0 +1,844 @@
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import AsyncGenerator, Generic, Set, Tuple, TypeVar
+
+from kiln_ai.adapters.chunkers.base_chunker import BaseChunker
+from kiln_ai.adapters.chunkers.chunker_registry import chunker_adapter_from_type
+from kiln_ai.adapters.embedding.base_embedding_adapter import BaseEmbeddingAdapter
+from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
+from kiln_ai.adapters.extractors.base_extractor import BaseExtractor, ExtractionInput
+from kiln_ai.adapters.extractors.extractor_registry import extractor_adapter_from_type
+from kiln_ai.adapters.rag.deduplication import (
+    deduplicate_chunk_embeddings,
+    deduplicate_chunked_documents,
+    deduplicate_extractions,
+    filter_documents_by_tags,
+)
+from kiln_ai.adapters.rag.progress import LogMessage, RagProgress
+from kiln_ai.adapters.vector_store.base_vector_store_adapter import (
+    DocumentWithChunksAndEmbeddings,
+)
+from kiln_ai.adapters.vector_store.vector_store_registry import (
+    vector_store_adapter_for_config,
+)
+from kiln_ai.datamodel import Project
+from kiln_ai.datamodel.basemodel import ID_TYPE, KilnAttachmentModel
+from kiln_ai.datamodel.chunk import Chunk, ChunkedDocument, ChunkerConfig
+from kiln_ai.datamodel.embedding import ChunkEmbeddings, Embedding, EmbeddingConfig
+from kiln_ai.datamodel.extraction import (
+    Document,
+    Extraction,
+    ExtractionSource,
+    ExtractorConfig,
+)
+from kiln_ai.datamodel.rag import RagConfig
+from kiln_ai.datamodel.vector_store import VectorStoreConfig
+from kiln_ai.utils.async_job_runner import AsyncJobRunner, AsyncJobRunnerObserver
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+from kiln_ai.utils.filesystem_cache import FilesystemCache
+from kiln_ai.utils.lock import shared_async_lock_manager
+from pydantic import BaseModel, ConfigDict, Field
+
+# We set the timeout high because current UX is likely to cause the user triggering
+# multiple RAG Workflows whose subconfigs (e.g. same extractor) may be shared and take
+# a long time to complete, causing whichever ones are waiting on the lock to time out
+# before they are likely to start.
+LOCK_TIMEOUT_SECONDS = 60 * 60  # 1 hour
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExtractorJob:
+    doc: Document
+    extractor_config: ExtractorConfig
+
+
+@dataclass
+class ChunkerJob:
+    extraction: Extraction
+    chunker_config: ChunkerConfig
+
+
+@dataclass
+class EmbeddingJob:
+    chunked_document: ChunkedDocument
+    embedding_config: EmbeddingConfig
+
+
+class RagStepRunnerProgress(BaseModel):
+    success_count: int | None = Field(
+        description="The number of items that have been processed",
+        default=None,
+    )
+    error_count: int | None = Field(
+        description="The number of items that have errored",
+        default=None,
+    )
+    logs: list[LogMessage] = Field(
+        description="A list of log messages to display to the user",
+        default_factory=list,
+    )
+
+
+T = TypeVar("T")
+
+
+class GenericErrorCollector(AsyncJobRunnerObserver[T], Generic[T]):
+    def __init__(
+        self,
+    ):
+        self.errors: list[Tuple[T, Exception]] = []
+
+    async def on_success(self, job: T):
+        pass
+
+    async def on_error(self, job: T, error: Exception):
+        self.errors.append((job, error))
+
+    def get_errors(
+        self,
+        start_idx: int = 0,
+    ) -> tuple[list[Tuple[T, Exception]], int]:
+        """Returns a tuple of: ((job, error), index of the last error)"""
+        if start_idx < 0:
+            raise ValueError("start_idx must be non-negative")
+        if start_idx >= len(self.errors):
+            return [], start_idx
+        if start_idx > 0:
+            return self.errors[start_idx : len(self.errors)], len(self.errors)
+        return self.errors, len(self.errors)
+
+    def get_error_count(self) -> int:
+        return len(self.errors)
+
+
+class RagWorkflowStepNames(str, Enum):
+    EXTRACTING = "extracting"
+    CHUNKING = "chunking"
+    EMBEDDING = "embedding"
+    INDEXING = "indexing"
+
+
+async def execute_extractor_job(job: ExtractorJob, extractor: BaseExtractor) -> bool:
+    if job.doc.path is None:
+        raise ValueError("Document path is not set")
+
+    output = await extractor.extract(
+        extraction_input=ExtractionInput(
+            path=job.doc.original_file.attachment.resolve_path(job.doc.path.parent),
+            mime_type=job.doc.original_file.mime_type,
+        )
+    )
+
+    extraction = Extraction(
+        parent=job.doc,
+        extractor_config_id=job.extractor_config.id,
+        output=KilnAttachmentModel.from_data(
+            data=output.content,
+            mime_type=output.content_format,
+        ),
+        source=ExtractionSource.PASSTHROUGH
+        if output.is_passthrough
+        else ExtractionSource.PROCESSED,
+    )
+    extraction.save_to_file()
+
+    return True
+
+
+async def execute_chunker_job(job: ChunkerJob, chunker: BaseChunker) -> bool:
+    extraction_output_content = await job.extraction.output_content()
+    if extraction_output_content is None:
+        raise ValueError("Extraction output content is not set")
+
+    chunking_result = await chunker.chunk(
+        extraction_output_content,
+    )
+    if chunking_result is None:
+        raise ValueError("Chunking result is not set")
+
+    chunked_document = ChunkedDocument(
+        parent=job.extraction,
+        chunker_config_id=job.chunker_config.id,
+        chunks=[
+            Chunk(
+                content=KilnAttachmentModel.from_data(
+                    data=chunk.text,
+                    mime_type="text/plain",
+                ),
+            )
+            for chunk in chunking_result.chunks
+        ],
+    )
+    chunked_document.save_to_file()
+    return True
+
+
+async def execute_embedding_job(
+    job: EmbeddingJob, embedding_adapter: BaseEmbeddingAdapter
+) -> bool:
+    chunks_text = await job.chunked_document.load_chunks_text()
+    if chunks_text is None or len(chunks_text) == 0:
+        raise ValueError(
+            f"Failed to load chunks for chunked document: {job.chunked_document.id}"
+        )
+
+    chunk_embedding_result = await embedding_adapter.generate_embeddings(
+        input_texts=chunks_text
+    )
+    if chunk_embedding_result is None:
+        raise ValueError(
+            f"Failed to generate embeddings for chunked document: {job.chunked_document.id}"
+        )
+
+    chunk_embeddings = ChunkEmbeddings(
+        parent=job.chunked_document,
+        embedding_config_id=job.embedding_config.id,
+        embeddings=[
+            Embedding(
+                vector=embedding.vector,
+            )
+            for embedding in chunk_embedding_result.embeddings
+        ],
+    )
+
+    chunk_embeddings.save_to_file()
+    return True
+
+
+class AbstractRagStepRunner(ABC):
+    @abstractmethod
+    def stage(self) -> RagWorkflowStepNames:
+        pass
+
+    # async keyword in the abstract prototype causes a type error in pyright
+    # so we need to remove it, but the concrete implementation should declare async
+    @abstractmethod
+    def run(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> AsyncGenerator[RagStepRunnerProgress, None]:
+        pass
+
+
+class RagExtractionStepRunner(AbstractRagStepRunner):
+    def __init__(
+        self,
+        project: Project,
+        extractor_config: ExtractorConfig,
+        concurrency: int = 10,
+        rag_config: RagConfig | None = None,
+        filesystem_cache: FilesystemCache | None = None,
+    ):
+        self.project = project
+        self.extractor_config = extractor_config
+        self.lock_key = f"docs:extract:{self.extractor_config.id}"
+        self.concurrency = concurrency
+        self.rag_config = rag_config
+        self.filesystem_cache = filesystem_cache
+
+    def stage(self) -> RagWorkflowStepNames:
+        return RagWorkflowStepNames.EXTRACTING
+
+    def has_extraction(self, document: Document, extractor_id: ID_TYPE) -> bool:
+        for ex in document.extractions(readonly=True):
+            if ex.extractor_config_id == extractor_id:
+                return True
+        return False
+
+    async def collect_jobs(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> list[ExtractorJob]:
+        jobs: list[ExtractorJob] = []
+        target_extractor_config_id = self.extractor_config.id
+
+        documents = self.project.documents(readonly=True)
+        if self.rag_config and self.rag_config.tags:
+            documents = filter_documents_by_tags(documents, self.rag_config.tags)
+
+        for document in documents:
+            if (
+                document_ids is not None
+                and len(document_ids) > 0
+                and document.id not in document_ids
+            ):
+                continue
+            if not self.has_extraction(document, target_extractor_config_id):
+                jobs.append(
+                    ExtractorJob(
+                        doc=document,
+                        extractor_config=self.extractor_config,
+                    )
+                )
+        return jobs
+
+    async def run(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> AsyncGenerator[RagStepRunnerProgress, None]:
+        async with shared_async_lock_manager.acquire(
+            self.lock_key, timeout=LOCK_TIMEOUT_SECONDS
+        ):
+            jobs = await self.collect_jobs(document_ids=document_ids)
+            extractor = extractor_adapter_from_type(
+                self.extractor_config.extractor_type,
+                self.extractor_config,
+                self.filesystem_cache,
+            )
+
+            observer = GenericErrorCollector()
+            runner = AsyncJobRunner(
+                jobs=jobs,
+                run_job_fn=lambda job: execute_extractor_job(job, extractor),
+                concurrency=self.concurrency,
+                observers=[observer],
+            )
+
+            error_idx = 0
+            async for progress in runner.run():
+                yield RagStepRunnerProgress(
+                    success_count=progress.complete,
+                    error_count=observer.get_error_count(),
+                )
+
+                # the errors are being accumulated in the observer so we need to flush them to the caller
+                if observer.get_error_count() > 0:
+                    errors, error_idx = observer.get_errors(error_idx)
+                    for job, error in errors:
+                        yield RagStepRunnerProgress(
+                            logs=[
+                                LogMessage(
+                                    level="error",
+                                    message=f"Error extracting document: {job.doc.path}: {error}",
+                                )
+                            ],
+                        )
+
+
+class RagChunkingStepRunner(AbstractRagStepRunner):
+    def __init__(
+        self,
+        project: Project,
+        extractor_config: ExtractorConfig,
+        chunker_config: ChunkerConfig,
+        concurrency: int = 10,
+        rag_config: RagConfig | None = None,
+    ):
+        self.project = project
+        self.extractor_config = extractor_config
+        self.chunker_config = chunker_config
+        self.lock_key = f"docs:chunk:{self.chunker_config.id}"
+        self.concurrency = concurrency
+        self.rag_config = rag_config
+
+    def stage(self) -> RagWorkflowStepNames:
+        return RagWorkflowStepNames.CHUNKING
+
+    def has_chunks(self, extraction: Extraction, chunker_id: ID_TYPE) -> bool:
+        for cd in extraction.chunked_documents(readonly=True):
+            if cd.chunker_config_id == chunker_id:
+                return True
+        return False
+
+    async def collect_jobs(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> list[ChunkerJob]:
+        target_extractor_config_id = self.extractor_config.id
+        target_chunker_config_id = self.chunker_config.id
+
+        jobs: list[ChunkerJob] = []
+        documents = self.project.documents(readonly=True)
+        if self.rag_config and self.rag_config.tags:
+            documents = filter_documents_by_tags(documents, self.rag_config.tags)
+
+        for document in documents:
+            if (
+                document_ids is not None
+                and len(document_ids) > 0
+                and document.id not in document_ids
+            ):
+                continue
+            for extraction in deduplicate_extractions(
+                document.extractions(readonly=True)
+            ):
+                if extraction.extractor_config_id == target_extractor_config_id:
+                    if not self.has_chunks(extraction, target_chunker_config_id):
+                        jobs.append(
+                            ChunkerJob(
+                                extraction=extraction,
+                                chunker_config=self.chunker_config,
+                            )
+                        )
+        return jobs
+
+    async def run(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> AsyncGenerator[RagStepRunnerProgress, None]:
+        async with shared_async_lock_manager.acquire(
+            self.lock_key, timeout=LOCK_TIMEOUT_SECONDS
+        ):
+            jobs = await self.collect_jobs(document_ids=document_ids)
+            chunker = chunker_adapter_from_type(
+                self.chunker_config.chunker_type,
+                self.chunker_config,
+            )
+            observer = GenericErrorCollector()
+            runner = AsyncJobRunner(
+                jobs=jobs,
+                run_job_fn=lambda job: execute_chunker_job(job, chunker),
+                concurrency=self.concurrency,
+                observers=[observer],
+            )
+
+            error_idx = 0
+            async for progress in runner.run():
+                yield RagStepRunnerProgress(
+                    success_count=progress.complete,
+                    error_count=observer.get_error_count(),
+                )
+
+                # the errors are being accumulated in the observer so we need to flush them to the caller
+                if observer.get_error_count() > 0:
+                    errors, error_idx = observer.get_errors(error_idx)
+                    for job, error in errors:
+                        yield RagStepRunnerProgress(
+                            logs=[
+                                LogMessage(
+                                    level="error",
+                                    message=f"Error chunking document: {job.extraction.path}: {error}",
+                                )
+                            ],
+                        )
+
+
+class RagEmbeddingStepRunner(AbstractRagStepRunner):
+    def __init__(
+        self,
+        project: Project,
+        extractor_config: ExtractorConfig,
+        chunker_config: ChunkerConfig,
+        embedding_config: EmbeddingConfig,
+        concurrency: int = 10,
+        rag_config: RagConfig | None = None,
+    ):
+        self.project = project
+        self.extractor_config = extractor_config
+        self.chunker_config = chunker_config
+        self.embedding_config = embedding_config
+        self.concurrency = concurrency
+        self.rag_config = rag_config
+        self.lock_key = f"docs:embedding:{self.embedding_config.id}"
+
+    def stage(self) -> RagWorkflowStepNames:
+        return RagWorkflowStepNames.EMBEDDING
+
+    def has_embeddings(self, chunked: ChunkedDocument, embedding_id: ID_TYPE) -> bool:
+        for emb in chunked.chunk_embeddings(readonly=True):
+            if emb.embedding_config_id == embedding_id:
+                return True
+        return False
+
+    async def collect_jobs(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> list[EmbeddingJob]:
+        target_extractor_config_id = self.extractor_config.id
+        target_chunker_config_id = self.chunker_config.id
+        target_embedding_config_id = self.embedding_config.id
+
+        jobs: list[EmbeddingJob] = []
+        documents = self.project.documents(readonly=True)
+        if self.rag_config and self.rag_config.tags:
+            documents = filter_documents_by_tags(documents, self.rag_config.tags)
+
+        for document in documents:
+            if (
+                document_ids is not None
+                and len(document_ids) > 0
+                and document.id not in document_ids
+            ):
+                continue
+            for extraction in deduplicate_extractions(
+                document.extractions(readonly=True)
+            ):
+                if extraction.extractor_config_id == target_extractor_config_id:
+                    for chunked_document in deduplicate_chunked_documents(
+                        extraction.chunked_documents(readonly=True)
+                    ):
+                        if (
+                            chunked_document.chunker_config_id
+                            == target_chunker_config_id
+                        ):
+                            if not self.has_embeddings(
+                                chunked_document, target_embedding_config_id
+                            ):
+                                jobs.append(
+                                    EmbeddingJob(
+                                        chunked_document=chunked_document,
+                                        embedding_config=self.embedding_config,
+                                    )
+                                )
+        return jobs
+
+    async def run(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> AsyncGenerator[RagStepRunnerProgress, None]:
+        async with shared_async_lock_manager.acquire(
+            self.lock_key, timeout=LOCK_TIMEOUT_SECONDS
+        ):
+            jobs = await self.collect_jobs(document_ids=document_ids)
+            embedding_adapter = embedding_adapter_from_type(
+                self.embedding_config,
+            )
+
+            observer = GenericErrorCollector()
+            runner = AsyncJobRunner(
+                jobs=jobs,
+                run_job_fn=lambda job: execute_embedding_job(job, embedding_adapter),
+                concurrency=self.concurrency,
+                observers=[observer],
+            )
+
+            error_idx = 0
+            async for progress in runner.run():
+                yield RagStepRunnerProgress(
+                    success_count=progress.complete,
+                    error_count=observer.get_error_count(),
+                )
+
+                # the errors are being accumulated in the observer so we need to flush them to the caller
+                if observer.get_error_count() > 0:
+                    errors, error_idx = observer.get_errors(error_idx)
+                    for job, error in errors:
+                        yield RagStepRunnerProgress(
+                            logs=[
+                                LogMessage(
+                                    level="error",
+                                    message=f"Error embedding document: {job.chunked_document.path}: {error}",
+                                )
+                            ],
+                        )
+
+
+class RagIndexingStepRunner(AbstractRagStepRunner):
+    def __init__(
+        self,
+        project: Project,
+        extractor_config: ExtractorConfig,
+        chunker_config: ChunkerConfig,
+        embedding_config: EmbeddingConfig,
+        vector_store_config: VectorStoreConfig,
+        rag_config: RagConfig,
+        concurrency: int = 10,
+        batch_size: int = 20,
+    ):
+        self.project = project
+        self.extractor_config = extractor_config
+        self.chunker_config = chunker_config
+        self.embedding_config = embedding_config
+        self.vector_store_config = vector_store_config
+        self.rag_config = rag_config
+        self.concurrency = concurrency
+        self.batch_size = batch_size
+
+    @property
+    def lock_key(self) -> str:
+        return f"rag:index:{self.vector_store_config.id}"
+
+    def stage(self) -> RagWorkflowStepNames:
+        return RagWorkflowStepNames.INDEXING
+
+    async def collect_records(
+        self,
+        batch_size: int,
+        document_ids: list[ID_TYPE] | None = None,
+    ) -> AsyncGenerator[list[DocumentWithChunksAndEmbeddings], None]:
+        target_extractor_config_id = self.extractor_config.id
+        target_chunker_config_id = self.chunker_config.id
+        target_embedding_config_id = self.embedding_config.id
+
+        # (document_id, chunked_document, embedding)
+        jobs: list[DocumentWithChunksAndEmbeddings] = []
+        documents = self.project.documents(readonly=True)
+        if self.rag_config and self.rag_config.tags:
+            documents = filter_documents_by_tags(documents, self.rag_config.tags)
+
+        for document in documents:
+            if (
+                document_ids is not None
+                and len(document_ids) > 0
+                and document.id not in document_ids
+            ):
+                continue
+            for extraction in deduplicate_extractions(
+                document.extractions(readonly=True)
+            ):
+                if extraction.extractor_config_id == target_extractor_config_id:
+                    for chunked_document in deduplicate_chunked_documents(
+                        extraction.chunked_documents(readonly=True)
+                    ):
+                        if (
+                            chunked_document.chunker_config_id
+                            == target_chunker_config_id
+                        ):
+                            for chunk_embedding in deduplicate_chunk_embeddings(
+                                chunked_document.chunk_embeddings(readonly=True)
+                            ):
+                                if (
+                                    chunk_embedding.embedding_config_id
+                                    == target_embedding_config_id
+                                ):
+                                    jobs.append(
+                                        DocumentWithChunksAndEmbeddings(
+                                            document_id=str(document.id),
+                                            chunked_document=chunked_document,
+                                            chunk_embeddings=chunk_embedding,
+                                        )
+                                    )
+
+                                    if len(jobs) >= batch_size:
+                                        yield jobs
+                                        jobs.clear()
+
+        if len(jobs) > 0:
+            yield jobs
+            jobs.clear()
+
+    async def count_total_chunks(self) -> int:
+        total_chunk_count = 0
+        async for documents in self.collect_records(batch_size=1):
+            total_chunk_count += len(documents[0].chunks)
+        return total_chunk_count
+
+    def get_all_target_document_ids(self) -> Set[str]:
+        documents = self.project.documents(readonly=True)
+        if self.rag_config and self.rag_config.tags:
+            documents = filter_documents_by_tags(documents, self.rag_config.tags)
+        return {str(document.id) for document in documents}
+
+    async def run(
+        self, document_ids: list[ID_TYPE] | None = None
+    ) -> AsyncGenerator[RagStepRunnerProgress, None]:
+        async with shared_async_lock_manager.acquire(
+            self.lock_key, timeout=LOCK_TIMEOUT_SECONDS
+        ):
+            vector_dimensions: int | None = None
+
+            # infer dimensionality - we peek into the first record to get the vector dimensions
+            # vector dimensions are not stored in the config because they are derived from the model
+            # and in some cases dynamic shortening of the vector (called Matryoshka Representation Learning)
+            async for doc_batch in self.collect_records(
+                batch_size=1,
+            ):
+                if len(doc_batch) == 0:
+                    # there are no records, because there may be nothing in the upstream steps at all yet
+                    return
+                else:
+                    doc = doc_batch[0]
+                    embedding = doc.embeddings[0]
+                    vector_dimensions = len(embedding.vector)
+                    break
+
+            if vector_dimensions is None:
+                raise ValueError("Vector dimensions are not set")
+
+            vector_store = await vector_store_adapter_for_config(
+                self.rag_config,
+                self.vector_store_config,
+            )
+
+            yield RagStepRunnerProgress(
+                success_count=0,
+                error_count=0,
+            )
+
+            async for doc_batch in self.collect_records(
+                batch_size=self.batch_size, document_ids=document_ids
+            ):
+                batch_chunk_count = 0
+                for doc in doc_batch:
+                    batch_chunk_count += len(doc.chunks)
+
+                try:
+                    await vector_store.add_chunks_with_embeddings(doc_batch)
+                    yield RagStepRunnerProgress(
+                        success_count=batch_chunk_count,
+                        error_count=0,
+                    )
+                except Exception as e:
+                    error_msg = f"Error indexing document batch starting with {doc_batch[0].document_id}: {e}"
+                    logger.error(error_msg, exc_info=True)
+                    yield RagStepRunnerProgress(
+                        success_count=0,
+                        error_count=batch_chunk_count,
+                        logs=[
+                            LogMessage(
+                                level="error",
+                                message=error_msg,
+                            ),
+                        ],
+                    )
+
+            # needed to reconcile and delete any chunks that are currently indexed but
+            # are no longer in our target set (because they were deleted or untagged)
+            await vector_store.delete_nodes_not_in_set(
+                self.get_all_target_document_ids()
+            )
+
+
+class RagWorkflowRunnerConfiguration(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    step_runners: list[AbstractRagStepRunner] = Field(
+        description="The step runners to run",
+    )
+
+    initial_progress: RagProgress = Field(
+        description="Initial progress state provided by the caller - progress will build on top of this",
+    )
+
+    rag_config: RagConfig = Field(
+        description="The rag config to use for the workflow",
+    )
+
+    extractor_config: ExtractorConfig = Field(
+        description="The extractor config to use for the workflow",
+    )
+
+    chunker_config: ChunkerConfig = Field(
+        description="The chunker config to use for the workflow",
+    )
+
+    embedding_config: EmbeddingConfig = Field(
+        description="The embedding config to use for the workflow",
+    )
+
+
+class RagWorkflowRunner:
+    def __init__(
+        self,
+        project: Project,
+        configuration: RagWorkflowRunnerConfiguration,
+    ):
+        self.project = project
+        self.configuration = configuration
+        self.step_runners: list[AbstractRagStepRunner] = configuration.step_runners
+        self.initial_progress = self.configuration.initial_progress
+        self.current_progress = self.initial_progress.model_copy()
+
+    @property
+    def lock_key(self) -> str:
+        return f"rag:run:{self.configuration.rag_config.id}"
+
+    def update_workflow_progress(
+        self, step_name: RagWorkflowStepNames, step_progress: RagStepRunnerProgress
+    ) -> RagProgress:
+        # merge the simpler step-specific progress with the broader RAG progress
+        match step_name:
+            case RagWorkflowStepNames.EXTRACTING:
+                if step_progress.success_count is not None:
+                    self.current_progress.total_document_extracted_count = max(
+                        self.current_progress.total_document_extracted_count,
+                        step_progress.success_count
+                        + self.initial_progress.total_document_extracted_count,
+                    )
+                if step_progress.error_count is not None:
+                    self.current_progress.total_document_extracted_error_count = max(
+                        self.current_progress.total_document_extracted_error_count,
+                        step_progress.error_count
+                        + self.initial_progress.total_document_extracted_error_count,
+                    )
+            case RagWorkflowStepNames.CHUNKING:
+                if step_progress.success_count is not None:
+                    self.current_progress.total_document_chunked_count = max(
+                        self.current_progress.total_document_chunked_count,
+                        step_progress.success_count
+                        + self.initial_progress.total_document_chunked_count,
+                    )
+                if step_progress.error_count is not None:
+                    self.current_progress.total_document_chunked_error_count = max(
+                        self.current_progress.total_document_chunked_error_count,
+                        step_progress.error_count
+                        + self.initial_progress.total_document_chunked_error_count,
+                    )
+            case RagWorkflowStepNames.EMBEDDING:
+                if step_progress.success_count is not None:
+                    self.current_progress.total_document_embedded_count = max(
+                        self.current_progress.total_document_embedded_count,
+                        step_progress.success_count
+                        + self.initial_progress.total_document_embedded_count,
+                    )
+                if step_progress.error_count is not None:
+                    self.current_progress.total_document_embedded_error_count = max(
+                        self.current_progress.total_document_embedded_error_count,
+                        step_progress.error_count
+                        + self.initial_progress.total_document_embedded_error_count,
+                    )
+            case RagWorkflowStepNames.INDEXING:
+                if step_progress.success_count is not None:
+                    self.current_progress.total_chunks_indexed_count += (
+                        step_progress.success_count
+                    )
+                if step_progress.error_count is not None:
+                    self.current_progress.total_chunks_indexed_error_count += (
+                        step_progress.error_count
+                    )
+            case _:
+                raise_exhaustive_enum_error(step_name)
+
+        self.current_progress.total_document_completed_count = min(
+            self.current_progress.total_document_extracted_count,
+            self.current_progress.total_document_chunked_count,
+            self.current_progress.total_document_embedded_count,
+        )
+
+        self.current_progress.total_chunk_completed_count = (
+            self.current_progress.total_chunks_indexed_count
+        )
+
+        self.current_progress.logs = step_progress.logs
+        return self.current_progress
+
+    async def run(
+        self,
+        stages_to_run: list[RagWorkflowStepNames] | None = None,
+        document_ids: list[ID_TYPE] | None = None,
+    ) -> AsyncGenerator[RagProgress, None]:
+        """
+        Runs the RAG workflow for the given stages and document ids.
+
+        :param stages_to_run: The stages to run. If None, all stages will be run.
+        :param document_ids: The document ids to run the workflow for. If None, all documents will be run.
+        """
+        yield self.initial_progress
+
+        async with shared_async_lock_manager.acquire(
+            self.lock_key, timeout=LOCK_TIMEOUT_SECONDS
+        ):
+            for step in self.step_runners:
+                if stages_to_run is not None and step.stage() not in stages_to_run:
+                    continue
+
+                # we need to know the total number of chunks to index to be able to
+                # calculate the progress on the client
+                if step.stage() == RagWorkflowStepNames.INDEXING and isinstance(
+                    step, RagIndexingStepRunner
+                ):
+                    self.current_progress.total_chunk_count = (
+                        await step.count_total_chunks()
+                    )
+                    # reset the indexing progress to 0 since we go through all the chunks again
+                    if not document_ids:
+                        self.initial_progress.total_chunks_indexed_count = 0
+                        self.current_progress.total_chunks_indexed_count = 0
+
+                    yield self.update_workflow_progress(
+                        step.stage(),
+                        RagStepRunnerProgress(
+                            success_count=0,
+                            error_count=0,
+                        ),
+                    )
+
+                async for progress in step.run(document_ids=document_ids):
+                    yield self.update_workflow_progress(step.stage(), progress)