two version of R2R are hereHEADmaster

1 files changed, 343 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/sdk/asnyc_methods/collections.py b/.venv/lib/python3.12/site-packages/sdk/asnyc_methods/collections.py
new file mode 100644
index 00000000..a768b72e
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/sdk/asnyc_methods/collections.py
@@ -0,0 +1,343 @@
+from typing import Any, Optional
+from uuid import UUID
+
+from shared.api.models import (
+    WrappedBooleanResponse,
+    WrappedCollectionResponse,
+    WrappedCollectionsResponse,
+    WrappedDocumentsResponse,
+    WrappedGenericMessageResponse,
+    WrappedUsersResponse,
+)
+
+
+class CollectionsSDK:
+    def __init__(self, client):
+        self.client = client
+
+    async def create(
+        self,
+        name: str,
+        description: Optional[str] = None,
+    ) -> WrappedCollectionResponse:
+        """Create a new collection.
+
+        Args:
+            name (str): Name of the collection
+            description (Optional[str]): Description of the collection
+
+        Returns:
+            WrappedCollectionResponse
+        """
+        data: dict[str, Any] = {"name": name, "description": description}
+        response_dict = await self.client._make_request(
+            "POST",
+            "collections",
+            json=data,
+            version="v3",
+        )
+
+        return WrappedCollectionResponse(**response_dict)
+
+    async def list(
+        self,
+        ids: Optional[list[str | UUID]] = None,
+        offset: Optional[int] = 0,
+        limit: Optional[int] = 100,
+    ) -> WrappedCollectionsResponse:
+        """List collections with pagination and filtering options.
+
+        Args:
+            ids (Optional[list[str | UUID]]): Filter collections by ids
+            offset (int, optional): Specifies the number of objects to skip. Defaults to 0.
+            limit (int, optional): Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.
+
+        Returns:
+            WrappedCollectionsResponse
+        """
+        params: dict = {
+            "offset": offset,
+            "limit": limit,
+        }
+        if ids:
+            params["ids"] = ids
+
+        response_dict = await self.client._make_request(
+            "GET", "collections", params=params, version="v3"
+        )
+
+        return WrappedCollectionsResponse(**response_dict)
+
+    async def retrieve(
+        self,
+        id: str | UUID,
+    ) -> WrappedCollectionResponse:
+        """Get detailed information about a specific collection.
+
+        Args:
+            id (str | UUID): Collection ID to retrieve
+
+        Returns:
+            WrappedCollectionResponse
+        """
+        response_dict = await self.client._make_request(
+            "GET", f"collections/{str(id)}", version="v3"
+        )
+
+        return WrappedCollectionResponse(**response_dict)
+
+    async def update(
+        self,
+        id: str | UUID,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        generate_description: Optional[bool] = False,
+    ) -> WrappedCollectionResponse:
+        """Update collection information.
+
+        Args:
+            id (str | UUID): Collection ID to update
+            name (Optional[str]): Optional new name for the collection
+            description (Optional[str]): Optional new description for the collection
+            generate_description (Optional[bool]): Whether to generate a new synthetic description for the collection.
+
+        Returns:
+            WrappedCollectionResponse
+        """
+        data: dict[str, Any] = {}
+        if name is not None:
+            data["name"] = name
+        if description is not None:
+            data["description"] = description
+        if generate_description:
+            data["generate_description"] = str(generate_description)
+
+        response_dict = await self.client._make_request(
+            "POST",
+            f"collections/{str(id)}",
+            json=data,
+            version="v3",
+        )
+
+        return WrappedCollectionResponse(**response_dict)
+
+    async def delete(
+        self,
+        id: str | UUID,
+    ) -> WrappedBooleanResponse:
+        """Delete a collection.
+
+        Args:
+            id (str | UUID): Collection ID to delete
+
+        Returns:
+            WrappedBooleanResponse
+        """
+        response_dict = await self.client._make_request(
+            "DELETE", f"collections/{str(id)}", version="v3"
+        )
+
+        return WrappedBooleanResponse(**response_dict)
+
+    async def list_documents(
+        self,
+        id: str | UUID,
+        offset: Optional[int] = 0,
+        limit: Optional[int] = 100,
+    ) -> WrappedDocumentsResponse:
+        """List all documents in a collection.
+
+        Args:
+            id (str | UUID): Collection ID
+            offset (int, optional): Specifies the number of objects to skip. Defaults to 0.
+            limit (int, optional): Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.
+
+        Returns:
+            WrappedDocumentsResponse
+        """
+        params: dict = {
+            "offset": offset,
+            "limit": limit,
+        }
+
+        response_dict = await self.client._make_request(
+            "GET",
+            f"collections/{str(id)}/documents",
+            params=params,
+            version="v3",
+        )
+
+        return WrappedDocumentsResponse(**response_dict)
+
+    async def add_document(
+        self,
+        id: str | UUID,
+        document_id: str | UUID,
+    ) -> WrappedGenericMessageResponse:
+        """Add a document to a collection.
+
+        Args:
+            id (str | UUID): Collection ID
+            document_id (str | UUID): Document ID to add
+
+        Returns:
+            WrappedGenericMessageResponse
+        """
+        response_dict = await self.client._make_request(
+            "POST",
+            f"collections/{str(id)}/documents/{str(document_id)}",
+            version="v3",
+        )
+
+        return WrappedGenericMessageResponse(**response_dict)
+
+    async def remove_document(
+        self,
+        id: str | UUID,
+        document_id: str | UUID,
+    ) -> WrappedBooleanResponse:
+        """Remove a document from a collection.
+
+        Args:
+            id (str | UUID): Collection ID
+            document_id (str | UUID): Document ID to remove
+
+        Returns:
+            WrappedBooleanResponse
+        """
+        response_dict = await self.client._make_request(
+            "DELETE",
+            f"collections/{str(id)}/documents/{str(document_id)}",
+            version="v3",
+        )
+
+        return WrappedBooleanResponse(**response_dict)
+
+    async def list_users(
+        self,
+        id: str | UUID,
+        offset: Optional[int] = 0,
+        limit: Optional[int] = 100,
+    ) -> WrappedUsersResponse:
+        """List all users in a collection.
+
+        Args:
+            id (str, UUID): Collection ID
+            offset (int, optional): Specifies the number of objects to skip. Defaults to 0.
+            limit (int, optional): Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.
+
+        Returns:
+            WrappedUsersResponse
+        """
+        params: dict = {
+            "offset": offset,
+            "limit": limit,
+        }
+
+        response_dict = await self.client._make_request(
+            "GET", f"collections/{str(id)}/users", params=params, version="v3"
+        )
+
+        return WrappedUsersResponse(**response_dict)
+
+    async def add_user(
+        self,
+        id: str | UUID,
+        user_id: str | UUID,
+    ) -> WrappedBooleanResponse:
+        """Add a user to a collection.
+
+        Args:
+            id (str | UUID): Collection ID
+            user_id (str | UUID): User ID to add
+
+        Returns:
+            WrappedBooleanResponse
+        """
+        response_dict = await self.client._make_request(
+            "POST", f"collections/{str(id)}/users/{str(user_id)}", version="v3"
+        )
+
+        return WrappedBooleanResponse(**response_dict)
+
+    async def remove_user(
+        self,
+        id: str | UUID,
+        user_id: str | UUID,
+    ) -> WrappedBooleanResponse:
+        """Remove a user from a collection.
+
+        Args:
+            id (str | UUID): Collection ID
+            user_id (str | UUID): User ID to remove
+
+        Returns:
+            WrappedBooleanResponse
+        """
+        response_dict = await self.client._make_request(
+            "DELETE",
+            f"collections/{str(id)}/users/{str(user_id)}",
+            version="v3",
+        )
+
+        return WrappedBooleanResponse(**response_dict)
+
+    async def extract(
+        self,
+        id: str | UUID,
+        settings: Optional[dict] = None,
+        run_with_orchestration: Optional[bool] = True,
+    ) -> WrappedGenericMessageResponse:
+        """Extract entities and relationships from documents in a collection.
+
+        Args:
+            id (str | UUID): Collection ID to extract from
+            settings (Optional[dict]): Settings for the entities and relationships extraction process
+            run_with_orchestration (Optional[bool]): Whether to run the extraction process with orchestration.
+                Defaults to True
+
+        Returns:
+            WrappedGenericMessageResponse
+        """
+        params = {"run_with_orchestration": run_with_orchestration}
+
+        data: dict[str, Any] = {}
+        if settings is not None:
+            data["settings"] = settings
+
+        response_dict = await self.client._make_request(
+            "POST",
+            f"collections/{str(id)}/extract",
+            params=params,
+            json=data or None,
+            version="v3",
+        )
+
+        return WrappedGenericMessageResponse(**response_dict)
+
+    async def retrieve_by_name(
+        self, name: str, owner_id: Optional[str] = None
+    ) -> WrappedCollectionResponse:
+        """Retrieve a collection by its name.
+
+        For non-superusers, the backend will use the authenticated user's ID.
+        For superusers, the caller must supply an owner_id to restrict the search.
+
+        Args:
+            name (str): The name of the collection to retrieve.
+            owner_id (Optional[str]): The owner ID to restrict the search. Required for superusers.
+
+        Returns:
+            WrappedCollectionResponse
+        """
+        query_params: dict[str, Any] = {}
+        if owner_id is not None:
+            query_params["owner_id"] = owner_id
+
+        response_dict = await self.client._make_request(
+            "GET",
+            f"collections/name/{name}",
+            params=query_params,
+            version="v3",
+        )
+        return WrappedCollectionResponse(**response_dict)