routes and deps

Author: nishika26

backend/app/api/deps.py

@@ -107,6 +107,30 @@ def get_current_user_org(
 CurrentUserOrg = Annotated[UserOrganization, Depends(get_current_user_org)]
 
 
+def get_current_user_org_project(
+    current_user: CurrentUser, session: SessionDep, request: Request
+) -> UserProjectOrg:
+    api_key = request.headers.get("X-API-KEY")
+    organization_id = None
+    project_id = None
+
+    if api_key:
+        api_key_record = get_api_key_by_value(session, api_key)
+        if api_key_record:
+            validate_organization(session, api_key_record.organization_id)
+            organization_id = api_key_record.organization_id
+            project_id = api_key_record.project_id
+
+    return UserProjectOrg(
+        **current_user.model_dump(),
+        organization_id=organization_id,
+        project_id=project_id,
+    )
+
+
+CurrentUserOrgproject = Annotated[UserProjectOrg, Depends(get_current_user_org_project)]
+
+
 def get_current_active_superuser(current_user: CurrentUser) -> User:
     if not current_user.is_superuser:
         raise HTTPException(

backend/app/api/routes/collections.py

@@ -11,7 +11,7 @@
 from pydantic import BaseModel, Field, HttpUrl
 from sqlalchemy.exc import NoResultFound, MultipleResultsFound, SQLAlchemyError
 
-from app.api.deps import CurrentUser, SessionDep
+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
@@ -169,65 +169,59 @@ def _backout(crud: OpenAIAssistantCrud, assistant_id: str):
 
 def do_create_collection(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgproject,
     request: CreationRequest,
     payload: ResponsePayload,
 ):
     client = OpenAI(api_key=settings.OPENAI_API_KEY)
-    if request.callback_url is None:
-        callback = SilentCallback(payload)
-    else:
-        callback = WebHookCallback(request.callback_url, payload)
-
-    #
-    # Create the assistant and vector store
-    #
+    callback = (
+        SilentCallback(payload)
+        if request.callback_url is None
+        else WebHookCallback(request.callback_url, payload)
+    )
 
     vector_store_crud = OpenAIVectorStoreCrud(client)
-    try:
-        vector_store = vector_store_crud.create()
-    except OpenAIError as err:
-        callback.fail(str(err))
-        return
-
+    assistant_crud = OpenAIAssistantCrud(client)
     storage = AmazonCloudStorage(current_user)
     document_crud = DocumentCrud(session, current_user.id)
-    assistant_crud = OpenAIAssistantCrud(client)
+    collection_crud = CollectionCrud(session, current_user.id)
 
-    docs = request(document_crud)
-    kwargs = dict(request.extract_super_type(AssistantOptions))
     try:
+        vector_store = vector_store_crud.create()
+
+        docs = request(document_crud)
         updates = vector_store_crud.update(vector_store.id, storage, docs)
         documents = list(updates)
+
+        kwargs = dict(request.extract_super_type(AssistantOptions))
         assistant = assistant_crud.create(vector_store.id, **kwargs)
-    except Exception as err:  # blanket to handle SQL and OpenAI errors
-        logging.error(f"File Search setup error: {err} ({type(err).__name__})")
-        vector_store_crud.delete(vector_store.id)
-        callback.fail(str(err))
-        return
 
-    #
-    # Store the results
-    #
+        # 3. Read and update collection with assistant info
+        collection = collection_crud.read_one(UUID(payload.key))
+        collection.llm_service_id = assistant.id
+        collection.llm_service_name = request.model
+        collection.status = "successfull"
+        collection.updated_at = now()
 
-    collection_crud = CollectionCrud(session, current_user.id)
-    collection = Collection(
-        id=UUID(payload.key),
-        llm_service_id=assistant.id,
-        llm_service_name=request.model,
-    )
-    try:
-        collection_crud.create(collection, documents)
-    except SQLAlchemyError as err:
-        _backout(assistant_crud, assistant.id)
-        callback.fail(str(err))
-        return
+        dc_crud = DocumentCollectionCrud(session)
+        dc_crud.create(collection, documents)
 
-    #
-    # Send back successful response
-    #
+        collection_crud._update(collection)
 
-    callback.success(collection.model_dump(mode="json"))
+        callback.success({"id": payload.key})
+    except Exception as err:
+        logging.error(f"[CollectionTask] Failed: {err} ({type(err).__name__})")
+
+        # 4. On failure, update collection status only
+        try:
+            collection = collection_crud.read_one(UUID(payload.key))
+            collection.status = "failed"
+            collection.updated_at = now()
+            collection_crud._update(collection)
+        except Exception as suberr:
+            logging.error(f"Failed to update collection status: {suberr}")
+
+        callback.fail(str(err))
 
 
 @router.post(
@@ -236,14 +230,27 @@ def do_create_collection(
 )
 def create_collection(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgproject,
     request: CreationRequest,
     background_tasks: BackgroundTasks,
 ):
     this = inspect.currentframe()
     route = router.url_path_for(this.f_code.co_name)
     payload = ResponsePayload("processing", route)
 
+    # 1. Create initial collection record
+    collection = Collection(
+        id=UUID(payload.key),
+        owner_id=current_user.id,
+        organization_id=current_user.organization_id,
+        project_id=current_user.project_id,
+        status="processing",
+    )
+
+    collection_crud = CollectionCrud(session, current_user.id)
+    collection_crud.create(collection, documents=[])
+
+    # 2. Launch background task
     background_tasks.add_task(
         do_create_collection,
         session,

backend/app/models/collection.py

@@ -1,27 +1,46 @@
 from uuid import UUID, uuid4
 from datetime import datetime
+from typing import Optional
 
 from sqlmodel import Field, Relationship, SQLModel
 
 from app.core.util import now
 from .user import User
+from .organization import Organization
+from .project import Project
+from enum import Enum
 
 
 class Collection(SQLModel, table=True):
-    id: UUID = Field(
-        default_factory=uuid4,
-        primary_key=True,
-    )
+    id: UUID = Field(default_factory=uuid4, primary_key=True)
+
     owner_id: int = Field(
         foreign_key="user.id",
         nullable=False,
         ondelete="CASCADE",
     )
-    llm_service_id: str
-    llm_service_name: str
-    created_at: datetime = Field(
-        default_factory=now,
+
+    organization_id: int = Field(
+        foreign_key="organization.id",
+        nullable=False,
+        ondelete="CASCADE",
+    )
+
+    project_id: int = Field(
+        foreign_key="project.id",
+        nullable=True,
+        ondelete="SET NULL",
     )
-    deleted_at: datetime | None
+
+    llm_service_id: Optional[str] = Field(default=None, nullable=True)
+    llm_service_name: Optional[str] = Field(default=None, nullable=True)
+
+    status: Optional[str] = None
+
+    created_at: datetime = Field(default_factory=now)
+    updated_at: datetime = Field(default_factory=now)
+    deleted_at: Optional[datetime] = None
 
     owner: User = Relationship(back_populates="collections")
+    organization: Organization = Relationship(back_populates="collections")
+    project: Project = Relationship(back_populates="collections")

backend/app/models/organization.py

@@ -10,6 +10,7 @@
     from .project import Project
     from .api_key import APIKey
     from .assistants import Assistant
+    from .collection import Collection
 
 
 # Shared properties for an Organization
@@ -48,6 +49,9 @@ class Organization(OrganizationBase, table=True):
     assistants: list["Assistant"] = Relationship(
         back_populates="organization", sa_relationship_kwargs={"cascade": "all, delete"}
     )
+    collections: list["Collection"] = Relationship(
+        back_populates="organization", sa_relationship_kwargs={"cascade": "all, delete"}
+    )
 
 
 # Properties to return via API

backend/app/models/project.py

@@ -44,6 +44,9 @@ class Project(ProjectBase, table=True):
         back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"}
     )
     organization: Optional["Organization"] = Relationship(back_populates="project")
+    collections: list["Collection"] = Relationship(
+        back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"}
+    )
 
 
 # Properties to return via API