arbi 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

arbi/resources/api/workspace.py

@@ -6,8 +6,8 @@ from typing import Optional
 
 import httpx
 
-from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
-from ..._utils import maybe_transform, async_maybe_transform
+from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
+from ..._utils import maybe_transform, strip_not_given, async_maybe_transform
 from ..._compat import cached_property
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._response import (
@@ -17,6 +17,7 @@ from ..._response import (
     async_to_streamed_response_wrapper,
 )
 from ...types.api import (
+    workspace_copy_params,
     workspace_share_params,
     workspace_update_params,
     workspace_remove_user_params,
@@ -24,6 +25,7 @@ from ...types.api import (
 )
 from ..._base_client import make_request_options
 from ...types.api.workspace_response import WorkspaceResponse
+from ...types.api.workspace_copy_response import WorkspaceCopyResponse
 from ...types.api.workspace_share_response import WorkspaceShareResponse
 from ...types.api.workspace_delete_response import WorkspaceDeleteResponse
 from ...types.api.workspace_get_tags_response import WorkspaceGetTagsResponse
@@ -64,6 +66,7 @@ class WorkspaceResource(SyncAPIResource):
         description: Optional[str] | Omit = omit,
         is_public: Optional[bool] | Omit = omit,
         name: Optional[str] | Omit = omit,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -76,7 +79,8 @@ class WorkspaceResource(SyncAPIResource):
         Changes
         are persisted to the database.
 
-        Only developers can change the is_public field.
+        Only developers can change the is_public field. When making a workspace public,
+        the backend uses the Workspace-Key header to get the workspace key.
 
         Args:
           extra_headers: Send extra headers
@@ -89,6 +93,7 @@ class WorkspaceResource(SyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._patch(
             f"/api/workspace/{workspace_ext_id}",
             body=maybe_transform(
@@ -142,6 +147,75 @@ class WorkspaceResource(SyncAPIResource):
             cast_to=WorkspaceDeleteResponse,
         )
 
+    def copy(
+        self,
+        workspace_ext_id: str,
+        *,
+        items: SequenceNotStr[str],
+        target_workspace_ext_id: str,
+        target_workspace_key: str | Omit = omit,
+        workspace_key: str | Omit = omit,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> WorkspaceCopyResponse:
+        """
+        Copy documents from source workspace to target workspace.
+
+        Requires:
+
+        - User must have access to source workspace (RLS enforced)
+        - Target workspace must exist and user must have access
+        - Workspace-Key header with source workspace key (optional for public
+          workspaces, required for private)
+        - Target-Workspace-Key header with target workspace key (required)
+
+        Copies:
+
+        - Document metadata (title, doc_date, shared status, etc.)
+        - MinIO encrypted files (downloaded to server memory, re-encrypted, uploaded)
+        - Qdrant vectors (with updated doc_ext_id and chunk_ext_id references)
+
+        Args:
+          items: List of document external IDs to copy (e.g., ['doc-a1b2c3d4', 'doc-e5f6g7h8'])
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+          timeout: Override the client-level default timeout for this request, in seconds
+        """
+        if not workspace_ext_id:
+            raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {
+            **strip_not_given(
+                {
+                    "target-workspace-key": target_workspace_key,
+                    "workspace-key": workspace_key,
+                }
+            ),
+            **(extra_headers or {}),
+        }
+        return self._post(
+            f"/api/workspace/{workspace_ext_id}/copy",
+            body=maybe_transform(
+                {
+                    "items": items,
+                    "target_workspace_ext_id": target_workspace_ext_id,
+                },
+                workspace_copy_params.WorkspaceCopyParams,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=WorkspaceCopyResponse,
+        )
+
     def create_protected(
         self,
         *,
@@ -160,6 +234,9 @@ class WorkspaceResource(SyncAPIResource):
         Sets up vector
         storage and associates the creator as the initial workspace user.
 
+        Server generates the workspace symmetric key and wraps it with the user's public
+        key. The wrapped key is returned in the response for client-side storage.
+
         Public workspaces are visible to all users and grant non-members limited access:
 
         - Non-members can view shared documents and tags
@@ -242,6 +319,7 @@ class WorkspaceResource(SyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -265,6 +343,7 @@ class WorkspaceResource(SyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._get(
             f"/api/workspace/{workspace_ext_id}/doctags",
             options=make_request_options(
@@ -277,6 +356,7 @@ class WorkspaceResource(SyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -300,6 +380,7 @@ class WorkspaceResource(SyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._get(
             f"/api/workspace/{workspace_ext_id}/documents",
             options=make_request_options(
@@ -320,7 +401,11 @@ class WorkspaceResource(SyncAPIResource):
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> WorkspaceGetStatsResponse:
         """
-        Retrieves conversation and document counts for a specific workspace.
+        Retrieves conversation and document counts with shared/private breakdown for a
+        specific workspace.
+
+        - Conversations are "shared" if they have at least one shared message
+        - Documents are "shared" if their shared field is True
 
         Args:
           extra_headers: Send extra headers
@@ -345,6 +430,7 @@ class WorkspaceResource(SyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -366,6 +452,7 @@ class WorkspaceResource(SyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._get(
             f"/api/workspace/{workspace_ext_id}/tags",
             options=make_request_options(
@@ -451,6 +538,7 @@ class WorkspaceResource(SyncAPIResource):
         workspace_ext_id: str,
         *,
         recipient_email: str,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -458,10 +546,11 @@ class WorkspaceResource(SyncAPIResource):
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> WorkspaceShareResponse:
-        """Share a workspace with another user via their email address.
+        """
+        Share a workspace with another user via their email address.
 
-        Securely transfers
-        workspace encryption keys to the recipient.
+        Client provides SealedBox-encrypted workspace key via Workspace-Key header.
+        Server decrypts it using session key, then wraps it with recipient's public key.
 
         Args:
           extra_headers: Send extra headers
@@ -474,6 +563,7 @@ class WorkspaceResource(SyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._post(
             f"/api/workspace/{workspace_ext_id}/share",
             body=maybe_transform({"recipient_email": recipient_email}, workspace_share_params.WorkspaceShareParams),
@@ -511,6 +601,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         description: Optional[str] | Omit = omit,
         is_public: Optional[bool] | Omit = omit,
         name: Optional[str] | Omit = omit,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -523,7 +614,8 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         Changes
         are persisted to the database.
 
-        Only developers can change the is_public field.
+        Only developers can change the is_public field. When making a workspace public,
+        the backend uses the Workspace-Key header to get the workspace key.
 
         Args:
           extra_headers: Send extra headers
@@ -536,6 +628,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._patch(
             f"/api/workspace/{workspace_ext_id}",
             body=await async_maybe_transform(
@@ -589,6 +682,75 @@ class AsyncWorkspaceResource(AsyncAPIResource):
             cast_to=WorkspaceDeleteResponse,
         )
 
+    async def copy(
+        self,
+        workspace_ext_id: str,
+        *,
+        items: SequenceNotStr[str],
+        target_workspace_ext_id: str,
+        target_workspace_key: str | Omit = omit,
+        workspace_key: str | Omit = omit,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> WorkspaceCopyResponse:
+        """
+        Copy documents from source workspace to target workspace.
+
+        Requires:
+
+        - User must have access to source workspace (RLS enforced)
+        - Target workspace must exist and user must have access
+        - Workspace-Key header with source workspace key (optional for public
+          workspaces, required for private)
+        - Target-Workspace-Key header with target workspace key (required)
+
+        Copies:
+
+        - Document metadata (title, doc_date, shared status, etc.)
+        - MinIO encrypted files (downloaded to server memory, re-encrypted, uploaded)
+        - Qdrant vectors (with updated doc_ext_id and chunk_ext_id references)
+
+        Args:
+          items: List of document external IDs to copy (e.g., ['doc-a1b2c3d4', 'doc-e5f6g7h8'])
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+          timeout: Override the client-level default timeout for this request, in seconds
+        """
+        if not workspace_ext_id:
+            raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {
+            **strip_not_given(
+                {
+                    "target-workspace-key": target_workspace_key,
+                    "workspace-key": workspace_key,
+                }
+            ),
+            **(extra_headers or {}),
+        }
+        return await self._post(
+            f"/api/workspace/{workspace_ext_id}/copy",
+            body=await async_maybe_transform(
+                {
+                    "items": items,
+                    "target_workspace_ext_id": target_workspace_ext_id,
+                },
+                workspace_copy_params.WorkspaceCopyParams,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=WorkspaceCopyResponse,
+        )
+
     async def create_protected(
         self,
         *,
@@ -607,6 +769,9 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         Sets up vector
         storage and associates the creator as the initial workspace user.
 
+        Server generates the workspace symmetric key and wraps it with the user's public
+        key. The wrapped key is returned in the response for client-side storage.
+
         Public workspaces are visible to all users and grant non-members limited access:
 
         - Non-members can view shared documents and tags
@@ -689,6 +854,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -712,6 +878,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._get(
             f"/api/workspace/{workspace_ext_id}/doctags",
             options=make_request_options(
@@ -724,6 +891,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -747,6 +915,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._get(
             f"/api/workspace/{workspace_ext_id}/documents",
             options=make_request_options(
@@ -767,7 +936,11 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> WorkspaceGetStatsResponse:
         """
-        Retrieves conversation and document counts for a specific workspace.
+        Retrieves conversation and document counts with shared/private breakdown for a
+        specific workspace.
+
+        - Conversations are "shared" if they have at least one shared message
+        - Documents are "shared" if their shared field is True
 
         Args:
           extra_headers: Send extra headers
@@ -792,6 +965,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         self,
         workspace_ext_id: str,
         *,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -813,6 +987,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._get(
             f"/api/workspace/{workspace_ext_id}/tags",
             options=make_request_options(
@@ -900,6 +1075,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         workspace_ext_id: str,
         *,
         recipient_email: str,
+        workspace_key: str | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -907,10 +1083,11 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> WorkspaceShareResponse:
-        """Share a workspace with another user via their email address.
+        """
+        Share a workspace with another user via their email address.
 
-        Securely transfers
-        workspace encryption keys to the recipient.
+        Client provides SealedBox-encrypted workspace key via Workspace-Key header.
+        Server decrypts it using session key, then wraps it with recipient's public key.
 
         Args:
           extra_headers: Send extra headers
@@ -923,6 +1100,7 @@ class AsyncWorkspaceResource(AsyncAPIResource):
         """
         if not workspace_ext_id:
             raise ValueError(f"Expected a non-empty value for `workspace_ext_id` but received {workspace_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._post(
             f"/api/workspace/{workspace_ext_id}/share",
             body=await async_maybe_transform(
@@ -945,6 +1123,9 @@ class WorkspaceResourceWithRawResponse:
         self.delete = to_raw_response_wrapper(
             workspace.delete,
         )
+        self.copy = to_raw_response_wrapper(
+            workspace.copy,
+        )
         self.create_protected = to_raw_response_wrapper(
             workspace.create_protected,
         )
@@ -984,6 +1165,9 @@ class AsyncWorkspaceResourceWithRawResponse:
         self.delete = async_to_raw_response_wrapper(
             workspace.delete,
         )
+        self.copy = async_to_raw_response_wrapper(
+            workspace.copy,
+        )
         self.create_protected = async_to_raw_response_wrapper(
             workspace.create_protected,
         )
@@ -1023,6 +1207,9 @@ class WorkspaceResourceWithStreamingResponse:
         self.delete = to_streamed_response_wrapper(
             workspace.delete,
         )
+        self.copy = to_streamed_response_wrapper(
+            workspace.copy,
+        )
         self.create_protected = to_streamed_response_wrapper(
             workspace.create_protected,
         )
@@ -1062,6 +1249,9 @@ class AsyncWorkspaceResourceWithStreamingResponse:
         self.delete = async_to_streamed_response_wrapper(
             workspace.delete,
         )
+        self.copy = async_to_streamed_response_wrapper(
+            workspace.copy,
+        )
         self.create_protected = async_to_streamed_response_wrapper(
             workspace.create_protected,
         )

arbi/types/api/__init__.py

@@ -2,35 +2,33 @@
 
 from __future__ import annotations
 
-from .token import Token as Token
 from .doc_response import DocResponse as DocResponse
-from .user_response import UserResponse as UserResponse
 from .embedder_config import EmbedderConfig as EmbedderConfig
 from .reranker_config import RerankerConfig as RerankerConfig
 from .query_llm_config import QueryLlmConfig as QueryLlmConfig
 from .retriever_config import RetrieverConfig as RetrieverConfig
-from .sso_login_params import SSOLoginParams as SSOLoginParams
 from .title_llm_config import TitleLlmConfig as TitleLlmConfig
-from .sso_invite_params import SSOInviteParams as SSOInviteParams
 from .tag_create_params import TagCreateParams as TagCreateParams
 from .tag_update_params import TagUpdateParams as TagUpdateParams
 from .user_login_params import UserLoginParams as UserLoginParams
-from .sso_login_response import SSOLoginResponse as SSOLoginResponse
+from .user_invite_params import UserInviteParams as UserInviteParams
 from .workspace_response import WorkspaceResponse as WorkspaceResponse
 from .parser_config_param import ParserConfigParam as ParserConfigParam
-from .sso_invite_response import SSOInviteResponse as SSOInviteResponse
 from .tag_create_response import TagCreateResponse as TagCreateResponse
 from .tag_delete_response import TagDeleteResponse as TagDeleteResponse
 from .tag_update_response import TagUpdateResponse as TagUpdateResponse
+from .user_login_response import UserLoginResponse as UserLoginResponse
 from .chunker_config_param import ChunkerConfigParam as ChunkerConfigParam
 from .config_create_params import ConfigCreateParams as ConfigCreateParams
 from .document_view_params import DocumentViewParams as DocumentViewParams
+from .user_invite_response import UserInviteResponse as UserInviteResponse
 from .user_logout_response import UserLogoutResponse as UserLogoutResponse
 from .user_register_params import UserRegisterParams as UserRegisterParams
 from .embedder_config_param import EmbedderConfigParam as EmbedderConfigParam
 from .model_citation_config import ModelCitationConfig as ModelCitationConfig
 from .reranker_config_param import RerankerConfigParam as RerankerConfigParam
 from .tag_get_docs_response import TagGetDocsResponse as TagGetDocsResponse
+from .workspace_copy_params import WorkspaceCopyParams as WorkspaceCopyParams
 from .assistant_query_params import AssistantQueryParams as AssistantQueryParams
 from .config_create_response import ConfigCreateResponse as ConfigCreateResponse
 from .config_delete_response import ConfigDeleteResponse as ConfigDeleteResponse
@@ -40,6 +38,7 @@ from .query_llm_config_param import QueryLlmConfigParam as QueryLlmConfigParam
 from .retriever_config_param import RetrieverConfigParam as RetrieverConfigParam
 from .title_llm_config_param import TitleLlmConfigParam as TitleLlmConfigParam
 from .workspace_share_params import WorkspaceShareParams as WorkspaceShareParams
+from .workspace_copy_response import WorkspaceCopyResponse as WorkspaceCopyResponse
 from .workspace_update_params import WorkspaceUpdateParams as WorkspaceUpdateParams
 from .config_retrieve_response import ConfigRetrieveResponse as ConfigRetrieveResponse
 from .document_delete_response import DocumentDeleteResponse as DocumentDeleteResponse
@@ -57,18 +56,21 @@ from .user_verify_email_response import UserVerifyEmailResponse as UserVerifyEma
 from .conversation_share_response import ConversationShareResponse as ConversationShareResponse
 from .model_citation_config_param import ModelCitationConfigParam as ModelCitationConfigParam
 from .tag_remove_from_docs_params import TagRemoveFromDocsParams as TagRemoveFromDocsParams
+from .user_change_password_params import UserChangePasswordParams as UserChangePasswordParams
 from .workspace_get_tags_response import WorkspaceGetTagsResponse as WorkspaceGetTagsResponse
 from .config_get_versions_response import ConfigGetVersionsResponse as ConfigGetVersionsResponse
 from .conversation_delete_response import ConversationDeleteResponse as ConversationDeleteResponse
 from .document_get_parsed_response import DocumentGetParsedResponse as DocumentGetParsedResponse
 from .health_check_models_response import HealthCheckModelsResponse as HealthCheckModelsResponse
-from .sso_rotate_passcode_response import SSORotatePasscodeResponse as SSORotatePasscodeResponse
+from .user_check_sso_status_params import UserCheckSSOStatusParams as UserCheckSSOStatusParams
 from .workspace_get_stats_response import WorkspaceGetStatsResponse as WorkspaceGetStatsResponse
 from .workspace_get_users_response import WorkspaceGetUsersResponse as WorkspaceGetUsersResponse
 from .workspace_remove_user_params import WorkspaceRemoveUserParams as WorkspaceRemoveUserParams
 from .tag_remove_from_docs_response import TagRemoveFromDocsResponse as TagRemoveFromDocsResponse
+from .user_change_password_response import UserChangePasswordResponse as UserChangePasswordResponse
 from .user_list_workspaces_response import UserListWorkspacesResponse as UserListWorkspacesResponse
 from .health_check_services_response import HealthCheckServicesResponse as HealthCheckServicesResponse
+from .user_check_sso_status_response import UserCheckSSOStatusResponse as UserCheckSSOStatusResponse
 from .workspace_get_doctags_response import WorkspaceGetDoctagsResponse as WorkspaceGetDoctagsResponse
 from .workspace_remove_user_response import WorkspaceRemoveUserResponse as WorkspaceRemoveUserResponse
 from .document_upload_from_url_params import DocumentUploadFromURLParams as DocumentUploadFromURLParams
@@ -81,6 +83,9 @@ from .conversation_update_title_response import ConversationUpdateTitleResponse
 from .document_date_extractor_llm_config import DocumentDateExtractorLlmConfig as DocumentDateExtractorLlmConfig
 from .conversation_delete_message_response import ConversationDeleteMessageResponse as ConversationDeleteMessageResponse
 from .workspace_get_conversations_response import WorkspaceGetConversationsResponse as WorkspaceGetConversationsResponse
+from .conversation_retrieve_message_response import (
+    ConversationRetrieveMessageResponse as ConversationRetrieveMessageResponse,
+)
 from .conversation_retrieve_threads_response import (
     ConversationRetrieveThreadsResponse as ConversationRetrieveThreadsResponse,
 )

arbi/types/api/assistant_query_params.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 from typing import Dict, Union, Iterable, Optional
-from typing_extensions import Literal, Required, TypeAlias, TypedDict
+from typing_extensions import Literal, Required, Annotated, TypeAlias, TypedDict
 
 from ..._types import SequenceNotStr
+from ..._utils import PropertyInfo
 from ..chunk_param import ChunkParam
 
 __all__ = [
@@ -15,6 +16,7 @@ __all__ = [
     "ToolsModelCitationToolToolResponses",
     "ToolsRetrievalChunkToolInput",
     "ToolsRetrievalFullContextToolInput",
+    "ToolsTraceTool",
 ]
 
 
@@ -29,6 +31,8 @@ class AssistantQueryParams(TypedDict, total=False):
 
     tools: Dict[str, Tools]
 
+    workspace_key: Annotated[str, PropertyInfo(alias="workspace-key")]
+
 
 class ToolsModelCitationToolToolResponses(TypedDict, total=False):
     chunk_ids: Required[SequenceNotStr[str]]
@@ -68,4 +72,20 @@ class ToolsRetrievalFullContextToolInput(TypedDict, total=False):
     tool_responses: Dict[str, Iterable[ChunkParam]]
 
 
-Tools: TypeAlias = Union[ToolsModelCitationTool, ToolsRetrievalChunkToolInput, ToolsRetrievalFullContextToolInput]
+class ToolsTraceTool(TypedDict, total=False):
+    description: str
+
+    duration_seconds: Optional[float]
+
+    name: Literal["trace"]
+
+    start_time: Optional[float]
+
+    steps: Iterable[Dict[str, object]]
+
+    trace_id: Optional[str]
+
+
+Tools: TypeAlias = Union[
+    ToolsModelCitationTool, ToolsRetrievalChunkToolInput, ToolsRetrievalFullContextToolInput, ToolsTraceTool
+]

arbi/types/api/assistant_retrieve_params.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 from typing import Dict, Union, Iterable, Optional
-from typing_extensions import Literal, Required, TypeAlias, TypedDict
+from typing_extensions import Literal, Required, Annotated, TypeAlias, TypedDict
 
 from ..._types import SequenceNotStr
+from ..._utils import PropertyInfo
 from ..chunk_param import ChunkParam
 
 __all__ = [
@@ -15,6 +16,7 @@ __all__ = [
     "ToolsModelCitationToolToolResponses",
     "ToolsRetrievalChunkToolInput",
     "ToolsRetrievalFullContextToolInput",
+    "ToolsTraceTool",
 ]
 
 
@@ -29,6 +31,8 @@ class AssistantRetrieveParams(TypedDict, total=False):
 
     tools: Dict[str, Tools]
 
+    workspace_key: Annotated[str, PropertyInfo(alias="workspace-key")]
+
 
 class ToolsModelCitationToolToolResponses(TypedDict, total=False):
     chunk_ids: Required[SequenceNotStr[str]]
@@ -68,4 +72,20 @@ class ToolsRetrievalFullContextToolInput(TypedDict, total=False):
     tool_responses: Dict[str, Iterable[ChunkParam]]
 
 
-Tools: TypeAlias = Union[ToolsModelCitationTool, ToolsRetrievalChunkToolInput, ToolsRetrievalFullContextToolInput]
+class ToolsTraceTool(TypedDict, total=False):
+    description: str
+
+    duration_seconds: Optional[float]
+
+    name: Literal["trace"]
+
+    start_time: Optional[float]
+
+    steps: Iterable[Dict[str, object]]
+
+    trace_id: Optional[str]
+
+
+Tools: TypeAlias = Union[
+    ToolsModelCitationTool, ToolsRetrievalChunkToolInput, ToolsRetrievalFullContextToolInput, ToolsTraceTool
+]