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

arbi/_streaming.py

@@ -57,9 +57,8 @@ class Stream(Generic[_T]):
         for sse in iterator:
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
-        # Ensure the entire stream is consumed
-        for _sse in iterator:
-            ...
+        # As we might not fully consume the response stream, we need to close it explicitly
+        response.close()
 
     def __enter__(self) -> Self:
         return self
@@ -121,9 +120,8 @@ class AsyncStream(Generic[_T]):
         async for sse in iterator:
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
-        # Ensure the entire stream is consumed
-        async for _sse in iterator:
-            ...
+        # As we might not fully consume the response stream, we need to close it explicitly
+        await response.aclose()
 
     async def __aenter__(self) -> Self:
         return self

arbi/_version.py

@@ -1,4 +1,4 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 __title__ = "arbi"
-__version__ = "0.2.1"  # x-release-please-version
+__version__ = "0.3.1"  # x-release-please-version

arbi/resources/api/__init__.py

@@ -8,14 +8,6 @@ from .api import (
     APIResourceWithStreamingResponse,
     AsyncAPIResourceWithStreamingResponse,
 )
-from .sso import (
-    SSOResource,
-    AsyncSSOResource,
-    SSOResourceWithRawResponse,
-    AsyncSSOResourceWithRawResponse,
-    SSOResourceWithStreamingResponse,
-    AsyncSSOResourceWithStreamingResponse,
-)
 from .tag import (
     TagResource,
     AsyncTagResource,
@@ -88,12 +80,6 @@ __all__ = [
     "AsyncUserResourceWithRawResponse",
     "UserResourceWithStreamingResponse",
     "AsyncUserResourceWithStreamingResponse",
-    "SSOResource",
-    "AsyncSSOResource",
-    "SSOResourceWithRawResponse",
-    "AsyncSSOResourceWithRawResponse",
-    "SSOResourceWithStreamingResponse",
-    "AsyncSSOResourceWithStreamingResponse",
     "WorkspaceResource",
     "AsyncWorkspaceResource",
     "WorkspaceResourceWithRawResponse",

arbi/resources/api/api.py

@@ -5,14 +5,6 @@ from __future__ import annotations
 import httpx
 
 from ... import _resource
-from .sso import (
-    SSOResource,
-    AsyncSSOResource,
-    SSOResourceWithRawResponse,
-    AsyncSSOResourceWithRawResponse,
-    SSOResourceWithStreamingResponse,
-    AsyncSSOResourceWithStreamingResponse,
-)
 from .tag import (
     TagResource,
     AsyncTagResource,
@@ -95,10 +87,6 @@ class APIResource(_resource.SyncAPIResource):
     def user(self) -> UserResource:
         return UserResource(self._client)
 
-    @cached_property
-    def sso(self) -> SSOResource:
-        return SSOResource(self._client)
-
     @cached_property
     def workspace(self) -> WorkspaceResource:
         return WorkspaceResource(self._client)
@@ -171,10 +159,6 @@ class AsyncAPIResource(_resource.AsyncAPIResource):
     def user(self) -> AsyncUserResource:
         return AsyncUserResource(self._client)
 
-    @cached_property
-    def sso(self) -> AsyncSSOResource:
-        return AsyncSSOResource(self._client)
-
     @cached_property
     def workspace(self) -> AsyncWorkspaceResource:
         return AsyncWorkspaceResource(self._client)
@@ -254,10 +238,6 @@ class APIResourceWithRawResponse:
     def user(self) -> UserResourceWithRawResponse:
         return UserResourceWithRawResponse(self._api.user)
 
-    @cached_property
-    def sso(self) -> SSOResourceWithRawResponse:
-        return SSOResourceWithRawResponse(self._api.sso)
-
     @cached_property
     def workspace(self) -> WorkspaceResourceWithRawResponse:
         return WorkspaceResourceWithRawResponse(self._api.workspace)
@@ -299,10 +279,6 @@ class AsyncAPIResourceWithRawResponse:
     def user(self) -> AsyncUserResourceWithRawResponse:
         return AsyncUserResourceWithRawResponse(self._api.user)
 
-    @cached_property
-    def sso(self) -> AsyncSSOResourceWithRawResponse:
-        return AsyncSSOResourceWithRawResponse(self._api.sso)
-
     @cached_property
     def workspace(self) -> AsyncWorkspaceResourceWithRawResponse:
         return AsyncWorkspaceResourceWithRawResponse(self._api.workspace)
@@ -344,10 +320,6 @@ class APIResourceWithStreamingResponse:
     def user(self) -> UserResourceWithStreamingResponse:
         return UserResourceWithStreamingResponse(self._api.user)
 
-    @cached_property
-    def sso(self) -> SSOResourceWithStreamingResponse:
-        return SSOResourceWithStreamingResponse(self._api.sso)
-
     @cached_property
     def workspace(self) -> WorkspaceResourceWithStreamingResponse:
         return WorkspaceResourceWithStreamingResponse(self._api.workspace)
@@ -389,10 +361,6 @@ class AsyncAPIResourceWithStreamingResponse:
     def user(self) -> AsyncUserResourceWithStreamingResponse:
         return AsyncUserResourceWithStreamingResponse(self._api.user)
 
-    @cached_property
-    def sso(self) -> AsyncSSOResourceWithStreamingResponse:
-        return AsyncSSOResourceWithStreamingResponse(self._api.sso)
-
     @cached_property
     def workspace(self) -> AsyncWorkspaceResourceWithStreamingResponse:
         return AsyncWorkspaceResourceWithStreamingResponse(self._api.workspace)

arbi/resources/api/assistant.py

@@ -7,7 +7,7 @@ from typing import Dict, Optional
 import httpx
 
 from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
-from ..._utils import maybe_transform, async_maybe_transform
+from ..._utils import maybe_transform, strip_not_given, async_maybe_transform
 from ..._compat import cached_property
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._response import (
@@ -50,6 +50,7 @@ class AssistantResource(SyncAPIResource):
         config_ext_id: Optional[str] | Omit = omit,
         parent_message_ext_id: Optional[str] | Omit = omit,
         tools: Dict[str, assistant_retrieve_params.Tools] | 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,
@@ -71,6 +72,7 @@ class AssistantResource(SyncAPIResource):
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._post(
             "/api/assistant/retrieve",
             body=maybe_transform(
@@ -97,6 +99,7 @@ class AssistantResource(SyncAPIResource):
         config_ext_id: Optional[str] | Omit = omit,
         parent_message_ext_id: Optional[str] | Omit = omit,
         tools: Dict[str, assistant_query_params.Tools] | 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,
@@ -118,6 +121,7 @@ class AssistantResource(SyncAPIResource):
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._post(
             "/api/assistant/query",
             body=maybe_transform(
@@ -165,6 +169,7 @@ class AsyncAssistantResource(AsyncAPIResource):
         config_ext_id: Optional[str] | Omit = omit,
         parent_message_ext_id: Optional[str] | Omit = omit,
         tools: Dict[str, assistant_retrieve_params.Tools] | 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,
@@ -186,6 +191,7 @@ class AsyncAssistantResource(AsyncAPIResource):
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._post(
             "/api/assistant/retrieve",
             body=await async_maybe_transform(
@@ -212,6 +218,7 @@ class AsyncAssistantResource(AsyncAPIResource):
         config_ext_id: Optional[str] | Omit = omit,
         parent_message_ext_id: Optional[str] | Omit = omit,
         tools: Dict[str, assistant_query_params.Tools] | 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,
@@ -233,6 +240,7 @@ class AsyncAssistantResource(AsyncAPIResource):
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._post(
             "/api/assistant/query",
             body=await async_maybe_transform(

arbi/resources/api/conversation/conversation.py

@@ -12,8 +12,8 @@ from .user import (
     UserResourceWithStreamingResponse,
     AsyncUserResourceWithStreamingResponse,
 )
-from ...._types import Body, Query, Headers, NotGiven, not_given
-from ...._utils import maybe_transform, async_maybe_transform
+from ...._types import Body, Omit, Query, Headers, NotGiven, 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 (
@@ -28,6 +28,7 @@ from ....types.api.conversation_share_response import ConversationShareResponse
 from ....types.api.conversation_delete_response import ConversationDeleteResponse
 from ....types.api.conversation_update_title_response import ConversationUpdateTitleResponse
 from ....types.api.conversation_delete_message_response import ConversationDeleteMessageResponse
+from ....types.api.conversation_retrieve_message_response import ConversationRetrieveMessageResponse
 from ....types.api.conversation_retrieve_threads_response import ConversationRetrieveThreadsResponse
 
 __all__ = ["ConversationResource", "AsyncConversationResource"]
@@ -128,6 +129,42 @@ class ConversationResource(SyncAPIResource):
             cast_to=ConversationDeleteMessageResponse,
         )
 
+    def retrieve_message(
+        self,
+        message_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,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> ConversationRetrieveMessageResponse:
+        """
+        Get a single message with full details including decrypted execution trace.
+        Always includes the trace with decrypted sensitive fields.
+
+        Args:
+          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 message_ext_id:
+            raise ValueError(f"Expected a non-empty value for `message_ext_id` but received {message_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
+        return self._get(
+            f"/api/conversation/message/{message_ext_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=ConversationRetrieveMessageResponse,
+        )
+
     def retrieve_threads(
         self,
         conversation_ext_id: str,
@@ -338,6 +375,42 @@ class AsyncConversationResource(AsyncAPIResource):
             cast_to=ConversationDeleteMessageResponse,
         )
 
+    async def retrieve_message(
+        self,
+        message_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,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> ConversationRetrieveMessageResponse:
+        """
+        Get a single message with full details including decrypted execution trace.
+        Always includes the trace with decrypted sensitive fields.
+
+        Args:
+          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 message_ext_id:
+            raise ValueError(f"Expected a non-empty value for `message_ext_id` but received {message_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
+        return await self._get(
+            f"/api/conversation/message/{message_ext_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=ConversationRetrieveMessageResponse,
+        )
+
     async def retrieve_threads(
         self,
         conversation_ext_id: str,
@@ -465,6 +538,9 @@ class ConversationResourceWithRawResponse:
         self.delete_message = to_raw_response_wrapper(
             conversation.delete_message,
         )
+        self.retrieve_message = to_raw_response_wrapper(
+            conversation.retrieve_message,
+        )
         self.retrieve_threads = to_raw_response_wrapper(
             conversation.retrieve_threads,
         )
@@ -490,6 +566,9 @@ class AsyncConversationResourceWithRawResponse:
         self.delete_message = async_to_raw_response_wrapper(
             conversation.delete_message,
         )
+        self.retrieve_message = async_to_raw_response_wrapper(
+            conversation.retrieve_message,
+        )
         self.retrieve_threads = async_to_raw_response_wrapper(
             conversation.retrieve_threads,
         )
@@ -515,6 +594,9 @@ class ConversationResourceWithStreamingResponse:
         self.delete_message = to_streamed_response_wrapper(
             conversation.delete_message,
         )
+        self.retrieve_message = to_streamed_response_wrapper(
+            conversation.retrieve_message,
+        )
         self.retrieve_threads = to_streamed_response_wrapper(
             conversation.retrieve_threads,
         )
@@ -540,6 +622,9 @@ class AsyncConversationResourceWithStreamingResponse:
         self.delete_message = async_to_streamed_response_wrapper(
             conversation.delete_message,
         )
+        self.retrieve_message = async_to_streamed_response_wrapper(
+            conversation.retrieve_message,
+        )
         self.retrieve_threads = async_to_streamed_response_wrapper(
             conversation.retrieve_threads,
         )

arbi/resources/api/document/annotation.py

@@ -7,7 +7,7 @@ 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 ...._utils import maybe_transform, strip_not_given, async_maybe_transform
 from ...._compat import cached_property
 from ...._resource import SyncAPIResource, AsyncAPIResource
 from ...._response import (
@@ -51,6 +51,7 @@ class AnnotationResource(SyncAPIResource):
         note: Optional[str] | Omit = omit,
         page_ref: Optional[int] | Omit = omit,
         tag_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,6 +77,7 @@ class AnnotationResource(SyncAPIResource):
         """
         if not doc_ext_id:
             raise ValueError(f"Expected a non-empty value for `doc_ext_id` but received {doc_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._post(
             f"/api/document/{doc_ext_id}/annotation",
             body=maybe_transform(
@@ -99,6 +101,7 @@ class AnnotationResource(SyncAPIResource):
         doc_ext_id: str,
         note: Optional[str] | Omit = omit,
         page_ref: Optional[int] | 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,
@@ -122,6 +125,7 @@ class AnnotationResource(SyncAPIResource):
             raise ValueError(f"Expected a non-empty value for `doc_ext_id` but received {doc_ext_id!r}")
         if not doctag_ext_id:
             raise ValueError(f"Expected a non-empty value for `doctag_ext_id` but received {doctag_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return self._patch(
             f"/api/document/{doc_ext_id}/annotation/{doctag_ext_id}",
             body=maybe_transform(
@@ -201,6 +205,7 @@ class AsyncAnnotationResource(AsyncAPIResource):
         note: Optional[str] | Omit = omit,
         page_ref: Optional[int] | Omit = omit,
         tag_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,
@@ -226,6 +231,7 @@ class AsyncAnnotationResource(AsyncAPIResource):
         """
         if not doc_ext_id:
             raise ValueError(f"Expected a non-empty value for `doc_ext_id` but received {doc_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._post(
             f"/api/document/{doc_ext_id}/annotation",
             body=await async_maybe_transform(
@@ -249,6 +255,7 @@ class AsyncAnnotationResource(AsyncAPIResource):
         doc_ext_id: str,
         note: Optional[str] | Omit = omit,
         page_ref: Optional[int] | 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,
@@ -272,6 +279,7 @@ class AsyncAnnotationResource(AsyncAPIResource):
             raise ValueError(f"Expected a non-empty value for `doc_ext_id` but received {doc_ext_id!r}")
         if not doctag_ext_id:
             raise ValueError(f"Expected a non-empty value for `doctag_ext_id` but received {doctag_ext_id!r}")
+        extra_headers = {**strip_not_given({"workspace-key": workspace_key}), **(extra_headers or {})}
         return await self._patch(
             f"/api/document/{doc_ext_id}/annotation/{doctag_ext_id}",
             body=await async_maybe_transform(