Collection: making the module provider agnostic #508

nishika26 · 2025-12-24T11:16:37Z

Summary

Target issue is #489

Checklist

Before submitting a pull request, please ensure that you mark these task.

Ran fastapi run --reload app/main.py or docker compose up in the repository root and test.
If you've fixed a bug or added code that is tested and has test cases.

Notes

Please add here if any other information is required for the reviewer.

Summary by CodeRabbit

Release Notes

New Features
- Introduced provider-based collection management system to support multiple LLM providers
- Collections now store provider-specific configuration data
- Added OpenAI provider with support for customizable collection parameters
Refactor
- Restructured collection creation and deletion workflows to use provider abstraction
- Extended public model exports for improved API accessibility

_{✏️ Tip: You can customize this high-level summary in your review settings.}

coderabbitai · 2025-12-24T11:16:43Z

📝 Walkthrough

Walkthrough

A new provider abstraction system is introduced for collection management, replacing direct provider-specific logic with a unified interface. Changes include database schema extensions (provider enum, collection_blob column), reorganized data models for request/response handling, and refactored service layer operations through a registry-based provider pattern supporting OpenAI backend.

Changes

Cohort / File(s)	Summary
Database Schema Migration `backend/app/alembic/versions/041_adding_blob_column_in_collection_table.py`	Adds PostgreSQL ENUM `provider_enum` with "openai" value; introduces `collection_blob` (JSONB) and `provider` columns to collection table; populates existing NULL providers to "openai"; enforces NOT NULL constraint post-population; updates `llm_service_name` comment.
Data Models - Collection Definitions `backend/app/models/collection/request.py`	Introduces `ProviderType` enum; adds `provider` and `collection_blob` fields to `Collection`; replaces document handling with `DocumentInput` type; introduces `CreateCollectionParams`, `CreationRequest`, `DeletionRequest`, `CallbackRequest`, and `ProviderOptions` types; includes document deduplication and provider value normalization.
Data Models - Collection Responses `backend/app/models/collection/response.py`	New module defining public response types: `CreateCollectionResult`, `CollectionIDPublic`, `CollectionPublic`, and `CollectionWithDocsPublic` with appropriate fields and inheritance structure.
Model Exports `backend/app/models/collection/__init__.py`, `backend/app/models/__init__.py`	Aggregates and re-exports collection-related types (`CreationRequest`, `DeletionRequest`, `ProviderType`, `CreateCollectionParams`, `CreateCollectionResult`, etc.) for simplified public API imports.
Collection Services `backend/app/services/collections/create_collection.py`, `backend/app/services/collections/delete_collection.py`	Replaces direct provider-specific logic with provider abstraction via `get_llm_provider`; removes `with_assistant` parameter from `execute_job`; delegates creation/deletion operations to provider interface; stores provider-specific metadata (`llm_service_id`, `llm_service_name`, `collection_blob`) in database.
Collection Helpers `backend/app/services/collections/helpers.py`	Introduces `get_service_name(provider: str)` helper function mapping providers to service names; removes `OPENAI_VECTOR_STORE` constant; updates logic to use helper function.
Provider Abstraction - Foundation `backend/app/services/collections/providers/base.py`, `backend/app/services/collections/providers/openai.py`	Introduces abstract `BaseProvider` class with `create`, `delete`, and `cleanup` methods; implements `OpenAIProvider` handling vector store creation, optional assistant creation, and resource cleanup.
Provider Abstraction - Registry `backend/app/services/collections/providers/registry.py`, `backend/app/services/collections/providers/__init__.py`	Introduces `LLMProvider` registry mapping provider names to provider classes; adds `get_llm_provider` factory function resolving provider class and retrieving credentials; exports public provider APIs.
Test Updates `backend/app/tests/api/routes/collections/test_collection_info.py`, `backend/app/tests/api/routes/collections/test_collection_list.py`, `backend/app/tests/utils/collection.py`	Updates test assertions to use `get_service_name("openai")` helper; sets `provider=ProviderType.OPENAI` on test Collection instances.

Sequence Diagrams

sequenceDiagram
    participant Client
    participant CollectionService as Collection Service
    participant Provider as LLMProvider
    participant OpenAIAPI as OpenAI API
    participant Database as Database
    
    Client->>CollectionService: execute_job(CreationRequest)
    activate CollectionService
    
    rect rgb(200, 220, 255)
        Note over CollectionService: Initialize Provider
        CollectionService->>Provider: get_llm_provider(provider="openai")
        Provider->>Provider: Lookup credentials via registry
        Provider->>OpenAIAPI: Initialize OpenAI client
        activate Provider
    end
    
    rect rgb(220, 240, 220)
        Note over CollectionService: Delegate Creation
        CollectionService->>Provider: provider.create(CreationRequest, storage, DocumentCrud)
        Provider->>OpenAIAPI: Create vector store from batched documents
        Provider->>OpenAIAPI: Optionally create assistant if model/instructions provided
        Provider-->>CollectionService: CreateCollectionResult (llm_service_id, llm_service_name, collection_blob)
        deactivate Provider
    end
    
    rect rgb(240, 220, 220)
        Note over CollectionService: Persist to Database
        CollectionService->>Database: Store Collection with provider, collection_blob, llm_service_id
        Database-->>CollectionService: Stored
    end
    
    alt Success
        CollectionService-->>Client: Job completed
    else Failure
        rect rgb(255, 200, 200)
            Note over CollectionService: Cleanup on Failure
            CollectionService->>Provider: provider.cleanup(result)
            Provider->>OpenAIAPI: Delete created vector store/assistant
        end
        CollectionService-->>Client: Job failed
    end
    
    deactivate CollectionService

sequenceDiagram
    participant Client
    participant CollectionService as Collection Service
    participant Provider as LLMProvider
    participant OpenAIAPI as OpenAI API
    participant Database as Database
    
    Client->>CollectionService: delete_collection(collection_id)
    activate CollectionService
    
    rect rgb(240, 240, 240)
        Note over CollectionService: Fetch Collection & Initialize Provider
        CollectionService->>Database: Fetch Collection (includes llm_service_name, provider)
        Database-->>CollectionService: Collection object
        CollectionService->>Provider: get_llm_provider(provider)
        activate Provider
    end
    
    rect rgb(220, 240, 220)
        Note over CollectionService: Delete External Resource
        CollectionService->>Provider: provider.delete(collection)
        alt llm_service_name != "openai vector store"
            Provider->>OpenAIAPI: Delete assistant
        else llm_service_name == "openai vector store"
            Provider->>OpenAIAPI: Delete vector store
        end
        Provider-->>CollectionService: Deleted
        deactivate Provider
    end
    
    rect rgb(240, 220, 220)
        Note over CollectionService: Remove from Database
        CollectionService->>Database: Delete collection record
        Database-->>CollectionService: Deleted
    end
    
    CollectionService-->>Client: Deletion complete
    deactivate CollectionService

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

Kaapi v1.0: Database Comments #476: Adds llm_service_id, llm_service_name, and timestamp metadata to collection schema; complements the provider and collection_blob additions in this PR.
Collection: Adding input parameter "provider" #502: Introduces explicit ProviderOptions and provider parameter handling in collection creation requests; aligns with the provider abstraction pattern introduced here.

Suggested reviewers

avirajsingh7
Prajna1999
kartpop

Poem

🐰 A provider's registry hops into place,
Abstract and clean, with elegant grace!
Vector stores dance with assistants alike,
Through OpenAI's tunnels—oh what a sight!
Collections now bloom with their provider's light. 🌸

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The pull request title clearly summarizes the main change: refactoring the collection module to be provider-agnostic by introducing a provider abstraction layer and registry system.
Docstring Coverage	✅ Passed	Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch enhancement/collection_provider_agnostic

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

backend/app/services/collections/create_collection.py (1)
269-270: Potential NameError if CreationRequest parsing fails.

If CreationRequest(**request) on line 156 raises an exception, creation_request is never assigned. The check on line 269 will then raise a NameError.
Proposed fix: initialize creation_request before try block or guard the check
+    creation_request = None
+
     try:
         creation_request = CreationRequest(**request)
         # ...

     except Exception as err:
         # ...

-        if creation_request and creation_request.callback_url and collection_job:
+        if creation_request is not None and creation_request.callback_url and collection_job:
             failure_payload = build_failure_payload(collection_job, str(err))
             send_callback(creation_request.callback_url, failure_payload)

🧹 Nitpick comments (11)

backend/app/services/collections/helpers.py (1)
17-25: Consider raising an error or logging for unknown providers.

Returning an empty string for unknown providers could lead to silent failures downstream. Consider logging a warning or raising a ValueError for unsupported providers to make debugging easier.
🔎 Suggested improvement
 def get_service_name(provider: str) -> str:
     """Get the collection service name for a provider."""
     names = {
         "openai": "openai vector store",
         #   "bedrock": "bedrock knowledge base",
         #  "gemini": "gemini file search store",
     }
-    return names.get(provider.lower(), "")
+    service_name = names.get(provider.lower())
+    if service_name is None:
+        logger.warning(f"[get_service_name] Unknown provider: {provider}")
+        return ""
+    return service_name
backend/app/services/collections/providers/base.py (3)
30-53: Docstring parameters don't match the method signature.

The docstring mentions batch_size, with_assistant, and assistant_options parameters that don't exist in the actual method signature. Also:

Line 48: "CreateCollectionresult" → "CreateCollectionResult" (typo)

Line 51: "kb_blob" → "collection_blob" (field name mismatch)

Line 53: error message says "execute method" but method is named "create"
Proposed fix
     @abstractmethod
     def create(
         self,
         collection_request: CreationRequest,
         storage: CloudStorage,
         document_crud: DocumentCrud,
     ) -> CreateCollectionResult:
         """Create collection with documents and optionally an assistant.

         Args:
-            collection_params: Collection parameters (name, description, chunking_params, etc.)
+            collection_request: Creation request containing collection params and options
             storage: Cloud storage instance for file access
             document_crud: DocumentCrud instance for fetching documents
-            batch_size: Number of documents to process per batch
-            with_assistant: Whether to create an assistant/agent
-            assistant_options: Options for assistant creation (provider-specific)

         Returns:
-            CreateCollectionresult containing:
+            CreateCollectionResult containing:
             - llm_service_id: ID of the created resource (vector store or assistant)
             - llm_service_name: Name of the service
-            - kb_blob: All collection params except documents
+            - collection_blob: All collection params except documents
         """
-        raise NotImplementedError("Providers must implement execute method")
+        raise NotImplementedError("Providers must implement create method")
55-65: Docstring Args don't match the method signature.

The docstring mentions llm_service_id and llm_service_name as parameters, but the actual signature only accepts collection: Collection.
Proposed fix
     @abstractmethod
     def delete(self, collection: Collection) -> None:
         """Delete remote resources associated with a collection.

         Called when a collection is being deleted and remote resources need to be cleaned up.

         Args:
-            llm_service_id: ID of the resource to delete
-            llm_service_name: Name of the service (determines resource type)
+            collection: The collection whose remote resources should be deleted
         """
         raise NotImplementedError("Providers must implement delete method")
67-76: Typo in docstring.

Line 74: "CreateCollectionresult" should be "CreateCollectionResult".
Proposed fix
-            collection_result: The CreateCollectionresult returned from execute, containing resource IDs
+            collection_result: The CreateCollectionResult returned from create, containing resource IDs
backend/app/services/collections/create_collection.py (1)
35-42: Unused with_assistant parameter.

The with_assistant parameter is accepted but never used in start_job. The assistant creation logic is now determined by checking model and instructions in the provider. Consider removing this parameter if it's no longer needed.
Proposed fix
 def start_job(
     db: Session,
     request: CreationRequest,
     project_id: int,
     collection_job_id: UUID,
-    with_assistant: bool,
     organization_id: int,
 ) -> str:
backend/app/services/collections/providers/openai.py (4)
2-2: Unused import: Any.

The Any type is imported but not used in this file.
Proposed fix
 import logging
-from typing import Any
 
 from openai import OpenAI
24-26: Redundant self.client assignment.

super().__init__(client) already assigns self.client = client in BaseProvider.__init__. The second assignment on line 26 is redundant.
Proposed fix
     def __init__(self, client: OpenAI):
         super().__init__(client)
-        self.client = client
62-65: Log messages reference wrong method name.

The log prefix says [OpenAIProvider.execute] but the method is named create. Per coding guidelines, log messages should be prefixed with the function name.
Proposed fix for all occurrences in create method
             logger.info(
-                "[OpenAIProvider.execute] Vector store created | "
+                "[OpenAIProvider.create] Vector store created | "
                 f"vector_store_id={vector_store.id}, batches={len(docs_batches)}"
             )
Apply similar changes to lines 93-95, 104-105, and 114-118.
60-60: Consider explicit loop for generator consumption.

Using list() to consume a generator whose result is discarded can be unclear. A for loop or collections.deque(maxlen=0) pattern would make intent clearer.
Proposed alternative
-            list(vector_store_crud.update(vector_store.id, storage, docs_batches))
+            for _ in vector_store_crud.update(vector_store.id, storage, docs_batches):
+                pass
backend/app/services/collections/providers/registry.py (1)
61-69: Unreachable else branch and logging format.

The else branch (lines 65-69) is unreachable because LLMProvider.get(provider) on line 47 already raises ValueError for unsupported providers. Also, the log message on line 67 should use square brackets per coding guidelines: [get_llm_provider].
Proposed fix: remove unreachable code or convert to assertion
     if provider == LLMProvider.OPENAI:
         if "api_key" not in credentials:
             raise ValueError("OpenAI credentials not configured for this project.")
         client = OpenAI(api_key=credentials["api_key"])
-    else:
-        logger.error(
-            f"[get_llm_provider] Unsupported provider type requested: {provider}"
-        )
-        raise ValueError(f"Provider '{provider}' is not supported.")
+    else:
+        # This branch is unreachable as LLMProvider.get validates the provider,
+        # but kept as defensive programming for future provider additions.
+        raise AssertionError(f"Unhandled provider: {provider}")

     return provider_class(client=client)
backend/app/models/collection/response.py (1)

20-29: Add provider field to CollectionPublic.

The Collection database model includes a provider field (ProviderType enum) that represents the LLM provider (e.g., "openai"). This field is missing from CollectionPublic and should be exposed in the response schema. Per learnings, provider and llm_service_name serve different purposes—provider indicates the LLM provider name while llm_service_name specifies the particular service from that provider. Exposing both fields provides complete information to API consumers about the collection's LLM configuration.

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 91941f9 and 946e7c7.

📒 Files selected for processing (15)

backend/app/alembic/versions/041_adding_blob_column_in_collection_table.py
backend/app/models/__init__.py
backend/app/models/collection/__init__.py
backend/app/models/collection/request.py
backend/app/models/collection/response.py
backend/app/services/collections/create_collection.py
backend/app/services/collections/delete_collection.py
backend/app/services/collections/helpers.py
backend/app/services/collections/providers/__init__.py
backend/app/services/collections/providers/base.py
backend/app/services/collections/providers/openai.py
backend/app/services/collections/providers/registry.py
backend/app/tests/api/routes/collections/test_collection_info.py
backend/app/tests/api/routes/collections/test_collection_list.py
backend/app/tests/utils/collection.py

🧰 Additional context used

📓 Path-based instructions (6)

backend/app/services/**/*.py