Add owner_id to extractors

backend/db/models.py

@@ -141,6 +141,11 @@ class Extractor(TimestampedModel):
         server_default="",
         comment="The name of the extractor.",
     )
+    owner_id = Column(
+        UUID(as_uuid=True),
+        nullable=False,
+        comment="Owner uuid.",
+    )
     schema = Column(
         JSONB,
         nullable=False,
@@ -168,3 +173,16 @@ class Extractor(TimestampedModel):
 
     def __repr__(self) -> str:
         return f"<Extractor(id={self.uuid}, description={self.description})>"
+
+
+def validate_extractor_owner(
+    session: Session, extractor_id: UUID, owner_id: UUID
+) -> Extractor:
+    """Validate the extractor id."""
+    extractor = (
+        session.query(Extractor).filter_by(uuid=extractor_id, owner_id=owner_id).first()
+    )
+    if extractor is None:
+        return False
+    else:
+        return True

backend/server/api/examples.py

@@ -2,11 +2,11 @@
 from typing import Any, List
 from uuid import UUID
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Cookie, Depends, HTTPException
 from sqlalchemy.orm import Session
 from typing_extensions import Annotated, TypedDict
 
-from db.models import Example, get_session
+from db.models import Example, get_session, validate_extractor_owner
 
 router = APIRouter(
     prefix="/examples",
@@ -36,8 +36,12 @@ def create(
     create_request: CreateExample,
     *,
     session: Session = Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> CreateExampleResponse:
     """Endpoint to create an example."""
+    if not validate_extractor_owner(session, create_request["extractor_id"], owner_id):
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
+
     instance = Example(
         extractor_id=create_request["extractor_id"],
         content=create_request["content"],
@@ -55,8 +59,11 @@ def list(
     limit: int = 10,
     offset: int = 0,
     session=Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> List[Any]:
     """Endpoint to get all examples."""
+    if not validate_extractor_owner(session, extractor_id, owner_id):
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
     return (
         session.query(Example)
         .filter(Example.extractor_id == extractor_id)
@@ -68,7 +75,12 @@ def list(
 
 
 @router.delete("/{uuid}")
-def delete(uuid: UUID, *, session: Session = Depends(get_session)) -> None:
+def delete(
+    uuid: UUID, *, session: Session = Depends(get_session), owner_id: UUID = Cookie(...)
+) -> None:
     """Endpoint to delete an example."""
-    session.query(Example).filter(Example.uuid == str(uuid)).delete()
+    extractor_id = session.query(Example).filter_by(uuid=str(uuid)).first().extractor_id
+    if not validate_extractor_owner(session, extractor_id, owner_id):
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
+    session.query(Example).filter_by(uuid=str(uuid)).delete()
     session.commit()

backend/server/api/extract.py

@@ -1,13 +1,14 @@
 from typing import Literal, Optional
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
+from fastapi import APIRouter, Cookie, Depends, File, Form, HTTPException, UploadFile
 from sqlalchemy.orm import Session
 from typing_extensions import Annotated
 
 from db.models import Extractor, get_session
 from extraction.parsing import parse_binary_input
 from server.extraction_runnable import ExtractResponse, extract_entire_document
+from server.models import DEFAULT_MODEL
 from server.retrieval import extract_from_content
 
 router = APIRouter(
@@ -24,8 +25,9 @@ async def extract_using_existing_extractor(
     text: Optional[str] = Form(None),
     mode: Literal["entire_document", "retrieval"] = Form("entire_document"),
     file: Optional[UploadFile] = File(None),
-    model_name: Optional[str] = Form("default"),
+    model_name: Optional[str] = Form(DEFAULT_MODEL),
     session: Session = Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> ExtractResponse:
     """Endpoint that is used with an existing extractor.
 
@@ -35,9 +37,13 @@ async def extract_using_existing_extractor(
     if text is None and file is None:
         raise HTTPException(status_code=422, detail="No text or file provided.")
 
-    extractor = session.query(Extractor).filter(Extractor.uuid == extractor_id).scalar()
+    extractor = (
+        session.query(Extractor)
+        .filter_by(uuid=extractor_id, owner_id=owner_id)
+        .scalar()
+    )
     if extractor is None:
-        raise HTTPException(status_code=404, detail="Extractor not found.")
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
 
     if text:
         text_ = text

backend/server/api/extractors.py

@@ -2,12 +2,12 @@
 from typing import Any, Dict, List
 from uuid import UUID, uuid4
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Cookie, Depends, HTTPException
 from pydantic import BaseModel, Field, validator
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
-from db.models import Extractor, SharedExtractors, get_session
+from db.models import Extractor, SharedExtractors, get_session, validate_extractor_owner
 from server.validators import validate_json_schema
 
 router = APIRouter(
@@ -60,6 +60,7 @@ def share(
     uuid: UUID,
     *,
     session: Session = Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> ShareExtractorResponse:
     """Endpoint to share an extractor.
 
@@ -73,6 +74,8 @@ def share(
     Returns:
         The UUID for the shared extractor.
     """
+    if not validate_extractor_owner(session, uuid, owner_id):
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
     # Check if the extractor is already shared
     shared_extractor = (
         session.query(SharedExtractors)
@@ -104,12 +107,16 @@ def share(
 
 @router.post("")
 def create(
-    create_request: CreateExtractor, *, session: Session = Depends(get_session)
+    create_request: CreateExtractor,
+    *,
+    session: Session = Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> CreateExtractorResponse:
     """Endpoint to create an extractor."""
 
     instance = Extractor(
         name=create_request.name,
+        owner_id=owner_id,
         schema=create_request.json_schema,
         description=create_request.description,
         instruction=create_request.instruction,
@@ -120,11 +127,15 @@ def create(
 
 
 @router.get("/{uuid}")
-def get(uuid: UUID, *, session: Session = Depends(get_session)) -> Dict[str, Any]:
+def get(
+    uuid: UUID, *, session: Session = Depends(get_session), owner_id: UUID = Cookie(...)
+) -> Dict[str, Any]:
     """Endpoint to get an extractor."""
-    extractor = session.query(Extractor).filter(Extractor.uuid == str(uuid)).scalar()
+    extractor = (
+        session.query(Extractor).filter_by(uuid=str(uuid), owner_id=owner_id).scalar()
+    )
     if extractor is None:
-        raise HTTPException(status_code=404, detail="Extractor not found.")
+        raise HTTPException(status_code=404, detail="Extractor not found for owner.")
     return {
         "uuid": extractor.uuid,
         "name": extractor.name,
@@ -140,13 +151,22 @@ def list(
     limit: int = 10,
     offset: int = 0,
     session=Depends(get_session),
+    owner_id: UUID = Cookie(...),
 ) -> List[Any]:
     """Endpoint to get all extractors."""
-    return session.query(Extractor).limit(limit).offset(offset).all()
+    return (
+        session.query(Extractor)
+        .filter_by(owner_id=owner_id)
+        .limit(limit)
+        .offset(offset)
+        .all()
+    )
 
 
 @router.delete("/{uuid}")
-def delete(uuid: UUID, *, session: Session = Depends(get_session)) -> None:
+def delete(
+    uuid: UUID, *, session: Session = Depends(get_session), owner_id: UUID = Cookie(...)
+) -> None:
     """Endpoint to delete an extractor."""
-    session.query(Extractor).filter(Extractor.uuid == str(uuid)).delete()
+    session.query(Extractor).filter_by(uuid=str(uuid), owner_id=owner_id).delete()
     session.commit()

backend/server/extraction_runnable.py

@@ -15,7 +15,7 @@
 
 from db.models import Example, Extractor
 from extraction.utils import update_json_schema
-from server.models import get_chunk_size, get_model
+from server.models import DEFAULT_MODEL, get_chunk_size, get_model
 from server.validators import validate_json_schema
 
 
@@ -188,7 +188,7 @@ async def extract_entire_document(
     text_splitter = TokenTextSplitter(
         chunk_size=get_chunk_size(model_name),
         chunk_overlap=20,
-        model_name=model_name,
+        model_name=DEFAULT_MODEL,
     )
     texts = text_splitter.split_text(content)
     extraction_requests = [

backend/tests/unit_tests/api/test_api_defining_extractors.py

@@ -8,7 +8,9 @@ async def test_extractors_api() -> None:
     """This will test a few of the extractors API endpoints."""
     # First verify that the database is empty
     async with get_async_client() as client:
-        response = await client.get("/extractors")
+        owner_id = str(uuid.uuid4())
+        cookies = {"owner_id": owner_id}
+        response = await client.get("/extractors", cookies=cookies)
         assert response.status_code == 200
         assert response.json() == []
 
@@ -18,22 +20,39 @@ async def test_extractors_api() -> None:
             "schema": {"type": "object"},
             "instruction": "Test Instruction",
         }
-        response = await client.post("/extractors", json=create_request)
+        response = await client.post(
+            "/extractors", json=create_request, cookies=cookies
+        )
         assert response.status_code == 200
 
         # Verify that the extractor was created
-        response = await client.get("/extractors")
+        response = await client.get("/extractors", cookies=cookies)
         assert response.status_code == 200
+        get_response = response.json()
+        assert len(get_response) == 1
+
+        # Check cookies
+        bad_cookies = {"owner_id": str(uuid.uuid4())}
+        bad_response = await client.get("/extractors", cookies=bad_cookies)
+        assert bad_response.status_code == 200
+        assert len(bad_response.json()) == 0
+
+        # Check we need cookie to delete
+        uuid_str = get_response[0]["uuid"]
+        _ = uuid.UUID(uuid_str)  # assert valid uuid
+        bad_response = await client.delete(
+            f"/extractors/{uuid_str}", cookies=bad_cookies
+        )
+        # Check extractor was not deleted
+        response = await client.get("/extractors", cookies=cookies)
         assert len(response.json()) == 1
 
         # Verify that we can delete an extractor
-        get_response = response.json()
-        uuid_str = get_response[0]["uuid"]
         _ = uuid.UUID(uuid_str)  # assert valid uuid
-        response = await client.delete(f"/extractors/{uuid_str}")
+        response = await client.delete(f"/extractors/{uuid_str}", cookies=cookies)
         assert response.status_code == 200
 
-        get_response = await client.get("/extractors")
+        get_response = await client.get("/extractors", cookies=cookies)
         assert get_response.status_code == 200
         assert get_response.json() == []
 
@@ -43,39 +62,53 @@ async def test_extractors_api() -> None:
             "schema": {"type": "object"},
             "instruction": "Test Instruction",
         }
-        response = await client.post("/extractors", json=create_request)
+        response = await client.post(
+            "/extractors", json=create_request, cookies=cookies
+        )
         assert response.status_code == 200
 
         # Verify that the extractor was created
-        response = await client.get("/extractors")
+        response = await client.get("/extractors", cookies=cookies)
         assert response.status_code == 200
         assert len(response.json()) == 1
 
         # Verify that we can delete an extractor
         get_response = response.json()
         uuid_str = get_response[0]["uuid"]
         _ = uuid.UUID(uuid_str)  # assert valid uuid
-        response = await client.delete(f"/extractors/{uuid_str}")
+        response = await client.delete(f"/extractors/{uuid_str}", cookies=cookies)
         assert response.status_code == 200
 
-        get_response = await client.get("/extractors")
+        get_response = await client.get("/extractors", cookies=cookies)
         assert get_response.status_code == 200
         assert get_response.json() == []
 
-        # Verify that we can create an extractor
+        # Verify that we can create an extractor, including other properties
+        owner_id = str(uuid.uuid4())
         create_request = {
+            "name": "my extractor",
             "description": "Test Description",
             "schema": {"type": "object"},
             "instruction": "Test Instruction",
         }
-        response = await client.post("/extractors", json=create_request)
+        response = await client.post(
+            "/extractors", json=create_request, cookies=cookies
+        )
+        extractor_uuid = response.json()["uuid"]
         assert response.status_code == 200
+        response = await client.get(f"/extractors/{extractor_uuid}", cookies=cookies)
+        response_data = response.json()
+        assert extractor_uuid == response_data["uuid"]
+        assert "my extractor" == response_data["name"]
+        assert "owner_id" not in response_data
 
 
 async def test_sharing_extractor() -> None:
     """Test sharing an extractor."""
     async with get_async_client() as client:
-        response = await client.get("/extractors")
+        owner_id = str(uuid.uuid4())
+        cookies = {"owner_id": owner_id}
+        response = await client.get("/extractors", cookies=cookies)
         assert response.status_code == 200
         assert response.json() == []
         # Verify that we can create an extractor
@@ -85,23 +118,32 @@ async def test_sharing_extractor() -> None:
             "schema": {"type": "object"},
             "instruction": "Test Instruction",
         }
-        response = await client.post("/extractors", json=create_request)
+        response = await client.post(
+            "/extractors", json=create_request, cookies=cookies
+        )
         assert response.status_code == 200
 
-        uuid = response.json()["uuid"]
+        uuid_str = response.json()["uuid"]
 
-        # Verify that the extractor was created
-        response = await client.post(f"/extractors/{uuid}/share")
+        # Generate a share uuid
+        response = await client.post(f"/extractors/{uuid_str}/share", cookies=cookies)
         assert response.status_code == 200
         assert "share_uuid" in response.json()
         share_uuid = response.json()["share_uuid"]
 
         # Test idempotency
-        response = await client.post(f"/extractors/{uuid}/share")
+        response = await client.post(f"/extractors/{uuid_str}/share", cookies=cookies)
         assert response.status_code == 200
         assert "share_uuid" in response.json()
         assert response.json()["share_uuid"] == share_uuid
 
+        # Check cookies
+        bad_cookies = {"owner_id": str(uuid.uuid4())}
+        response = await client.post(
+            f"/extractors/{uuid_str}/share", cookies=bad_cookies
+        )
+        assert response.status_code == 404
+
         # Check that we can retrieve the shared extractor
         response = await client.get(f"/s/{share_uuid}")
         assert response.status_code == 200