Collection: Standardized OpenAI Creds Fetching

.env.example

@@ -43,7 +43,7 @@ DOCKER_IMAGE_FRONTEND=frontend
 AWS_ACCESS_KEY_ID=
 AWS_SECRET_ACCESS_KEY=
 AWS_DEFAULT_REGION=ap-south-1
-AWS_S3_BUCKET_PREFIX = "bucket-prefix-name"
+AWS_S3_BUCKET_PREFIX="bucket-prefix-name"
 
 # OpenAI
 

backend/app/api/routes/collections.py

@@ -15,8 +15,13 @@
 from app.api.deps import CurrentUser, SessionDep, CurrentUserOrgProject
 from app.core.cloud import AmazonCloudStorage
 from app.core.config import settings
-from app.core.util import now, raise_from_unknown, post_callback
-from app.crud import DocumentCrud, CollectionCrud, DocumentCollectionCrud
+from app.core.util import now, raise_from_unknown, post_callback, configure_openai
+from app.crud import (
+    DocumentCrud,
+    CollectionCrud,
+    DocumentCollectionCrud,
+    get_provider_credential,
+)
 from app.crud.rag import OpenAIVectorStoreCrud, OpenAIAssistantCrud
 from app.models import Collection, Document
 from app.models.collection import CollectionStatus
@@ -178,14 +183,27 @@
         )
 
 
+def mark_collection_failed(session, user_id, collection_id, reason: str):
+    try:
+        collection = CollectionCrud(session, user_id).read_one(collection_id)
+        collection.status = CollectionStatus.failed
+        collection.updated_at = now()
+        CollectionCrud(session, user_id)._update(collection)
+    except Exception as suberr:
+        logger.warning(
+            f"[do_create_collection] Failed to mark collection failed | {{'collection_id': '{collection_id}', 'reason': '{str(suberr)}'}}"
+        )
+
+
 def do_create_collection(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     request: CreationRequest,
     payload: ResponsePayload,
+    client,
 ):
     start_time = time.time()
-    client = OpenAI(api_key=settings.OPENAI_API_KEY)
+
     callback = (
         SilentCallback(payload)
         if request.callback_url is None
@@ -226,7 +244,7 @@
         collection_crud._update(collection)
 
         elapsed = time.time() - start_time
-        logging.info(
+        logger.info(
             f"[do_create_collection] Collection created: {collection.id} | Time: {elapsed:.2f}s | "
             f"Files: {len(flat_docs)} | Sizes: {file_sizes_kb} KB | Types: {list(file_exts)}"
         )
@@ -261,6 +279,19 @@
     request: CreationRequest,
     background_tasks: BackgroundTasks,
 ):
+    credentials = get_provider_credential(
+        session=session,
+        org_id=current_user.organization_id,
+        provider="openai",
+        project_id=current_user.project_id,
+    )
+    client, success = configure_openai(credentials)
+    if not success:
+        logger.error(
+            f"[create_collection] OpenAI API key not configured for org_id={current_user.organization_id}, project_id={current_user.project_id}"
+        )
+        raise HTTPException(status_code=400, detail="OpenAI is not configured")
+
     this = inspect.currentframe()
     route = router.url_path_for(this.f_code.co_name)
     payload = ResponsePayload("processing", route)
@@ -278,11 +309,7 @@
 
     # 2. Launch background task
     background_tasks.add_task(
-        do_create_collection,
-        session,
-        current_user,
-        request,
-        payload,
+        do_create_collection, session, current_user, request, payload, client
     )
 
     logger.info(

backend/app/api/routes/documents.py

@@ -3,13 +3,14 @@
 from typing import List
 from pathlib import Path
 
-from fastapi import APIRouter, File, UploadFile, Query
+from fastapi import APIRouter, File, UploadFile, Query, HTTPException
 from fastapi import Path as FastPath
 
-from app.crud import DocumentCrud, CollectionCrud
+from app.crud import DocumentCrud, CollectionCrud, get_provider_credential
 from app.models import Document
+from app.core.util import configure_openai
 from app.utils import APIResponse, load_description
-from app.api.deps import CurrentUser, SessionDep
+from app.api.deps import CurrentUser, SessionDep, CurrentUserOrgProject
 from app.core.cloud import AmazonCloudStorage
 from app.crud.rag import OpenAIAssistantCrud
 
@@ -65,10 +66,23 @@
 )
 def remove_doc(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     doc_id: UUID = FastPath(description="Document to delete"),
 ):
-    a_crud = OpenAIAssistantCrud()
+    credentials = get_provider_credential(
+        session=session,
+        org_id=current_user.organization_id,
+        provider="openai",
+        project_id=current_user.project_id,
+    )
+    client, success = configure_openai(credentials)
+    if not success:
+        logger.error(
+            f"[remove_doc] OpenAI API key not configured for org_id={current_user.organization_id}, project_id={current_user.project_id}"
+        )
+        raise HTTPException(status_code=400, detail="OpenAI is not configured")
+
+    a_crud = OpenAIAssistantCrud(client)
     d_crud = DocumentCrud(session, current_user.id)
     c_crud = CollectionCrud(session, current_user.id)
 
@@ -84,10 +98,23 @@
 )
 def permanent_delete_doc(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     doc_id: UUID = FastPath(description="Document to permanently delete"),
 ):
-    a_crud = OpenAIAssistantCrud()
+    credentials = get_provider_credential(
+        session=session,
+        org_id=current_user.organization_id,
+        provider="openai",
+        project_id=current_user.project_id,
+    )
+    client, success = configure_openai(credentials)
+    if not success:
+        logger.error(
+            f"[permanent_delete_doc] OpenAI API key not configured for org_id={current_user.organization_id}, project_id={current_user.project_id}"
+        )
+        raise HTTPException(status_code=400, detail="OpenAI is not configured")
+
+    a_crud = OpenAIAssistantCrud(client)
     d_crud = DocumentCrud(session, current_user.id)
     c_crud = CollectionCrud(session, current_user.id)
     storage = AmazonCloudStorage(current_user)

backend/app/crud/__init__.py

@@ -37,6 +37,7 @@
     get_key_by_org,
     update_creds_for_org,
     remove_creds_for_org,
+    get_provider_credential,
 )
 
 from .thread_results import upsert_thread_result, get_thread_result

backend/app/crud/rag/open_ai.py

@@ -90,8 +90,10 @@
 
 
 class OpenAICrud:
-    def __init__(self, client=None):
-        self.client = client or OpenAI(api_key=settings.OPENAI_API_KEY)
+    def __init__(self, client):
+        if client is None:
+            raise ValueError("OpenAI client is not configured")
+        self.client = client
 
 
 class OpenAIVectorStoreCrud(OpenAICrud):

backend/app/tests/api/routes/collections/test_create_collections.py

@@ -3,18 +3,18 @@
 import io
 
 import openai_responses
-from sqlmodel import Session, select
+from sqlmodel import Session
 from fastapi.testclient import TestClient
-from openai import OpenAIError
+from unittest.mock import MagicMock, patch
 
 from app.core.config import settings
 from app.tests.utils.document import DocumentStore
-from app.tests.utils.utils import openai_credentials, get_user_from_api_key
+from app.tests.utils.utils import get_user_from_api_key
 from app.main import app
 from app.crud.collection import CollectionCrud
-from app.api.routes.collections import CreationRequest, ResponsePayload
 from app.seed_data.seed_data import seed_database
 from app.models.collection import CollectionStatus
+from app.tests.utils.collections_openai_mock import get_mock_openai_client
 
 client = TestClient(app)
 
@@ -51,17 +51,21 @@ def head_object(self, Bucket, Key):
     monkeypatch.setattr("boto3.client", lambda service: FakeS3Client())
 
 
-@pytest.mark.usefixtures("openai_credentials")
+@patch("app.api.routes.collections.configure_openai")
+@patch("app.api.routes.collections.get_provider_credential")
 class TestCollectionRouteCreate:
     _n_documents = 5
 
-    @openai_responses.mock()
     def test_create_collection_success(
         self,
+        mock_get_credential,
+        mock_configure_openai,
         client: TestClient,
         db: Session,
+        api_key_headers: dict[str, str],
     ):
-        store = DocumentStore(db)
+        # Setup test documents
+        store = DocumentStore(db, api_key_headers)
         documents = store.fill(self._n_documents)
         doc_ids = [str(doc.id) for doc in documents]
 
@@ -73,9 +77,13 @@ def test_create_collection_success(
             "temperature": 0.1,
         }
         original_api_key = "ApiKey No3x47A5qoIGhm0kVKjQ77dhCqEdWRIQZlEPzzzh7i8"
-
         headers = {"X-API-KEY": original_api_key}
 
+        mock_get_credential.return_value = {"api_key": "test_api_key"}
+
+        mock_openai_client = get_mock_openai_client()
+        mock_configure_openai.return_value = (mock_openai_client, True)
+
         response = client.post(
             f"{settings.API_V1_STR}/collections/create",
             json=body,
@@ -89,8 +97,8 @@ def test_create_collection_success(
         assert metadata["status"] == CollectionStatus.processing.value
         assert UUID(metadata["key"])
 
+        # Confirm collection metadata in DB
         collection_id = UUID(metadata["key"])
-
         user = get_user_from_api_key(db, headers)
         collection = CollectionCrud(db, user.user_id).read_one(collection_id)
 

backend/app/tests/api/routes/documents/test_route_document_info.py

@@ -23,8 +23,9 @@ def test_response_is_success(
         db: Session,
         route: Route,
         crawler: WebCrawler,
+        api_key_headers: dict[str, str],
     ):
-        store = DocumentStore(db)
+        store = DocumentStore(db, api_key_headers)
         response = crawler.get(route.append(store.put()))
 
         assert response.is_success
@@ -34,8 +35,9 @@ def test_info_reflects_database(
         db: Session,
         route: Route,
         crawler: WebCrawler,
+        api_key_headers: dict[str, str],
     ):
-        store = DocumentStore(db)
+        store = DocumentStore(db, api_key_headers)
         document = store.put()
         source = DocumentComparator(document)
 
@@ -44,13 +46,10 @@ def test_info_reflects_database(
         assert source == target.data
 
     def test_cannot_info_unknown_document(
-        self,
-        db: Session,
-        route: Route,
-        crawler: Route,
+        self, db: Session, route: Route, crawler: Route, api_key_headers: dict[str, str]
     ):
         DocumentStore.clear(db)
-        maker = DocumentMaker(db)
+        maker = DocumentMaker(db, api_key_headers)
         response = crawler.get(route.append(next(maker)))
 
         assert response.is_error

backend/app/tests/api/routes/documents/test_route_document_list.py

@@ -47,8 +47,9 @@ def test_item_reflects_database(
         db: Session,
         route: QueryRoute,
         crawler: WebCrawler,
+        api_key_headers: dict[str, str],
     ):
-        store = DocumentStore(db)
+        store = DocumentStore(db, api_key_headers)
         source = DocumentComparator(store.put())
 
         response = httpx_to_standard(crawler.get(route))
@@ -77,8 +78,9 @@ def test_skip_greater_than_limit_is_difference(
         db: Session,
         route: QueryRoute,
         crawler: WebCrawler,
+        api_key_headers: dict[str, str],
     ):
-        store = DocumentStore(db)
+        store = DocumentStore(db, api_key_headers)
         limit = len(store.fill(self._ndocs))
         skip = limit // 2