vellum-ai 1.1.2__py3-none-any.whl → 1.1.3__py3-none-any.whl

vellum/client/resources/ad_hoc/client.py

@@ -12,14 +12,6 @@ from ...types.function_definition import FunctionDefinition
 from ...types.ad_hoc_expand_meta import AdHocExpandMeta
 from ...core.request_options import RequestOptions
 from ...types.ad_hoc_execute_prompt_event import AdHocExecutePromptEvent
-from ...core.serialization import convert_and_respect_annotation_metadata
-from ...core.pydantic_utilities import parse_obj_as
-import json
-from ...errors.bad_request_error import BadRequestError
-from ...errors.forbidden_error import ForbiddenError
-from ...errors.internal_server_error import InternalServerError
-from json.decoder import JSONDecodeError
-from ...core.api_error import ApiError
 from ...core.client_wrapper import AsyncClientWrapper
 from .raw_client import AsyncRawAdHocClient
 
@@ -239,93 +231,18 @@ class AdHocClient:
         for chunk in response:
             yield chunk
         """
-        with self._raw_client._client_wrapper.httpx_client.stream(
-            "v1/ad-hoc/execute-prompt-stream",
-            base_url=self._raw_client._client_wrapper.get_environment().predict,
-            method="POST",
-            json={
-                "ml_model": ml_model,
-                "input_values": convert_and_respect_annotation_metadata(
-                    object_=input_values, annotation=typing.Sequence[PromptRequestInput], direction="write"
-                ),
-                "input_variables": convert_and_respect_annotation_metadata(
-                    object_=input_variables, annotation=typing.Sequence[VellumVariable], direction="write"
-                ),
-                "parameters": convert_and_respect_annotation_metadata(
-                    object_=parameters, annotation=PromptParameters, direction="write"
-                ),
-                "settings": convert_and_respect_annotation_metadata(
-                    object_=settings, annotation=typing.Optional[PromptSettings], direction="write"
-                ),
-                "blocks": convert_and_respect_annotation_metadata(
-                    object_=blocks, annotation=typing.Sequence[PromptBlock], direction="write"
-                ),
-                "functions": convert_and_respect_annotation_metadata(
-                    object_=functions,
-                    annotation=typing.Optional[typing.Sequence[FunctionDefinition]],
-                    direction="write",
-                ),
-                "expand_meta": convert_and_respect_annotation_metadata(
-                    object_=expand_meta, annotation=typing.Optional[AdHocExpandMeta], direction="write"
-                ),
-            },
-            headers={
-                "content-type": "application/json",
-            },
+        with self._raw_client.adhoc_execute_prompt_stream(
+            ml_model=ml_model,
+            input_values=input_values,
+            input_variables=input_variables,
+            parameters=parameters,
+            blocks=blocks,
+            settings=settings,
+            functions=functions,
+            expand_meta=expand_meta,
             request_options=request_options,
-            omit=OMIT,
-        ) as _response:
-            try:
-                if 200 <= _response.status_code < 300:
-                    for _text in _response.iter_lines():
-                        try:
-                            if len(_text) == 0:
-                                continue
-                            yield typing.cast(
-                                AdHocExecutePromptEvent,
-                                parse_obj_as(
-                                    type_=AdHocExecutePromptEvent,  # type: ignore
-                                    object_=json.loads(_text),
-                                ),
-                            )
-                        except Exception:
-                            pass
-                    return
-                _response.read()
-                if _response.status_code == 400:
-                    raise BadRequestError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                if _response.status_code == 403:
-                    raise ForbiddenError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                if _response.status_code == 500:
-                    raise InternalServerError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                _response_json = _response.json()
-            except JSONDecodeError:
-                raise ApiError(status_code=_response.status_code, body=_response.text)
-            raise ApiError(status_code=_response.status_code, body=_response_json)
+        ) as r:
+            yield from r.data
 
 
 class AsyncAdHocClient:
@@ -552,90 +469,16 @@ class AsyncAdHocClient:
 
         asyncio.run(main())
         """
-        async with self._raw_client._client_wrapper.httpx_client.stream(
-            "v1/ad-hoc/execute-prompt-stream",
-            base_url=self._raw_client._client_wrapper.get_environment().predict,
-            method="POST",
-            json={
-                "ml_model": ml_model,
-                "input_values": convert_and_respect_annotation_metadata(
-                    object_=input_values, annotation=typing.Sequence[PromptRequestInput], direction="write"
-                ),
-                "input_variables": convert_and_respect_annotation_metadata(
-                    object_=input_variables, annotation=typing.Sequence[VellumVariable], direction="write"
-                ),
-                "parameters": convert_and_respect_annotation_metadata(
-                    object_=parameters, annotation=PromptParameters, direction="write"
-                ),
-                "settings": convert_and_respect_annotation_metadata(
-                    object_=settings, annotation=typing.Optional[PromptSettings], direction="write"
-                ),
-                "blocks": convert_and_respect_annotation_metadata(
-                    object_=blocks, annotation=typing.Sequence[PromptBlock], direction="write"
-                ),
-                "functions": convert_and_respect_annotation_metadata(
-                    object_=functions,
-                    annotation=typing.Optional[typing.Sequence[FunctionDefinition]],
-                    direction="write",
-                ),
-                "expand_meta": convert_and_respect_annotation_metadata(
-                    object_=expand_meta, annotation=typing.Optional[AdHocExpandMeta], direction="write"
-                ),
-            },
-            headers={
-                "content-type": "application/json",
-            },
+        async with self._raw_client.adhoc_execute_prompt_stream(
+            ml_model=ml_model,
+            input_values=input_values,
+            input_variables=input_variables,
+            parameters=parameters,
+            blocks=blocks,
+            settings=settings,
+            functions=functions,
+            expand_meta=expand_meta,
             request_options=request_options,
-            omit=OMIT,
-        ) as _response:
-            try:
-                if 200 <= _response.status_code < 300:
-                    async for _text in _response.aiter_lines():
-                        try:
-                            if len(_text) == 0:
-                                continue
-                            yield typing.cast(
-                                AdHocExecutePromptEvent,
-                                parse_obj_as(
-                                    type_=AdHocExecutePromptEvent,  # type: ignore
-                                    object_=json.loads(_text),
-                                ),
-                            )
-                        except Exception:
-                            pass
-                    return
-                await _response.aread()
-                if _response.status_code == 400:
-                    raise BadRequestError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                if _response.status_code == 403:
-                    raise ForbiddenError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                if _response.status_code == 500:
-                    raise InternalServerError(
-                        typing.cast(
-                            typing.Optional[typing.Any],
-                            parse_obj_as(
-                                type_=typing.Optional[typing.Any],  # type: ignore
-                                object_=_response.json(),
-                            ),
-                        )
-                    )
-                _response_json = _response.json()
-            except JSONDecodeError:
-                raise ApiError(status_code=_response.status_code, body=_response.text)
-            raise ApiError(status_code=_response.status_code, body=_response_json)
+        ) as r:
+            async for data in r.data:
+                yield data

vellum/client/resources/ad_hoc/raw_client.py

@@ -19,6 +19,8 @@ from ...errors.forbidden_error import ForbiddenError
 from ...errors.internal_server_error import InternalServerError
 from json.decoder import JSONDecodeError
 from ...core.api_error import ApiError
+import json
+import contextlib
 from ...core.client_wrapper import AsyncClientWrapper
 from ...core.http_response import AsyncHttpResponse
 
@@ -151,6 +153,143 @@ class RawAdHocClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
+    @contextlib.contextmanager
+    def adhoc_execute_prompt_stream(
+        self,
+        *,
+        ml_model: str,
+        input_values: typing.Sequence[PromptRequestInput],
+        input_variables: typing.Sequence[VellumVariable],
+        parameters: PromptParameters,
+        blocks: typing.Sequence[PromptBlock],
+        settings: typing.Optional[PromptSettings] = OMIT,
+        functions: typing.Optional[typing.Sequence[FunctionDefinition]] = OMIT,
+        expand_meta: typing.Optional[AdHocExpandMeta] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.Iterator[HttpResponse[typing.Iterator[AdHocExecutePromptEvent]]]:
+        """
+        Parameters
+        ----------
+        ml_model : str
+
+        input_values : typing.Sequence[PromptRequestInput]
+
+        input_variables : typing.Sequence[VellumVariable]
+
+        parameters : PromptParameters
+
+        blocks : typing.Sequence[PromptBlock]
+
+        settings : typing.Optional[PromptSettings]
+
+        functions : typing.Optional[typing.Sequence[FunctionDefinition]]
+
+        expand_meta : typing.Optional[AdHocExpandMeta]
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Yields
+        ------
+        typing.Iterator[HttpResponse[typing.Iterator[AdHocExecutePromptEvent]]]
+
+        """
+        with self._client_wrapper.httpx_client.stream(
+            "v1/ad-hoc/execute-prompt-stream",
+            base_url=self._client_wrapper.get_environment().predict,
+            method="POST",
+            json={
+                "ml_model": ml_model,
+                "input_values": convert_and_respect_annotation_metadata(
+                    object_=input_values, annotation=typing.Sequence[PromptRequestInput], direction="write"
+                ),
+                "input_variables": convert_and_respect_annotation_metadata(
+                    object_=input_variables, annotation=typing.Sequence[VellumVariable], direction="write"
+                ),
+                "parameters": convert_and_respect_annotation_metadata(
+                    object_=parameters, annotation=PromptParameters, direction="write"
+                ),
+                "settings": convert_and_respect_annotation_metadata(
+                    object_=settings, annotation=typing.Optional[PromptSettings], direction="write"
+                ),
+                "blocks": convert_and_respect_annotation_metadata(
+                    object_=blocks, annotation=typing.Sequence[PromptBlock], direction="write"
+                ),
+                "functions": convert_and_respect_annotation_metadata(
+                    object_=functions,
+                    annotation=typing.Optional[typing.Sequence[FunctionDefinition]],
+                    direction="write",
+                ),
+                "expand_meta": convert_and_respect_annotation_metadata(
+                    object_=expand_meta, annotation=typing.Optional[AdHocExpandMeta], direction="write"
+                ),
+            },
+            headers={
+                "content-type": "application/json",
+            },
+            request_options=request_options,
+            omit=OMIT,
+        ) as _response:
+
+            def stream() -> HttpResponse[typing.Iterator[AdHocExecutePromptEvent]]:
+                try:
+                    if 200 <= _response.status_code < 300:
+
+                        def _iter():
+                            for _text in _response.iter_lines():
+                                try:
+                                    if len(_text) == 0:
+                                        continue
+                                    yield typing.cast(
+                                        AdHocExecutePromptEvent,
+                                        parse_obj_as(
+                                            type_=AdHocExecutePromptEvent,  # type: ignore
+                                            object_=json.loads(_text),
+                                        ),
+                                    )
+                                except Exception:
+                                    pass
+                            return
+
+                        return HttpResponse(response=_response, data=_iter())
+                    _response.read()
+                    if _response.status_code == 400:
+                        raise BadRequestError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    if _response.status_code == 403:
+                        raise ForbiddenError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    if _response.status_code == 500:
+                        raise InternalServerError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    _response_json = _response.json()
+                except JSONDecodeError:
+                    raise ApiError(status_code=_response.status_code, body=_response.text)
+                raise ApiError(status_code=_response.status_code, body=_response_json)
+
+            yield stream()
+
 
 class AsyncRawAdHocClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
@@ -276,3 +415,140 @@ class AsyncRawAdHocClient:
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
+
+    @contextlib.asynccontextmanager
+    async def adhoc_execute_prompt_stream(
+        self,
+        *,
+        ml_model: str,
+        input_values: typing.Sequence[PromptRequestInput],
+        input_variables: typing.Sequence[VellumVariable],
+        parameters: PromptParameters,
+        blocks: typing.Sequence[PromptBlock],
+        settings: typing.Optional[PromptSettings] = OMIT,
+        functions: typing.Optional[typing.Sequence[FunctionDefinition]] = OMIT,
+        expand_meta: typing.Optional[AdHocExpandMeta] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[AdHocExecutePromptEvent]]]:
+        """
+        Parameters
+        ----------
+        ml_model : str
+
+        input_values : typing.Sequence[PromptRequestInput]
+
+        input_variables : typing.Sequence[VellumVariable]
+
+        parameters : PromptParameters
+
+        blocks : typing.Sequence[PromptBlock]
+
+        settings : typing.Optional[PromptSettings]
+
+        functions : typing.Optional[typing.Sequence[FunctionDefinition]]
+
+        expand_meta : typing.Optional[AdHocExpandMeta]
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Yields
+        ------
+        typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[AdHocExecutePromptEvent]]]
+
+        """
+        async with self._client_wrapper.httpx_client.stream(
+            "v1/ad-hoc/execute-prompt-stream",
+            base_url=self._client_wrapper.get_environment().predict,
+            method="POST",
+            json={
+                "ml_model": ml_model,
+                "input_values": convert_and_respect_annotation_metadata(
+                    object_=input_values, annotation=typing.Sequence[PromptRequestInput], direction="write"
+                ),
+                "input_variables": convert_and_respect_annotation_metadata(
+                    object_=input_variables, annotation=typing.Sequence[VellumVariable], direction="write"
+                ),
+                "parameters": convert_and_respect_annotation_metadata(
+                    object_=parameters, annotation=PromptParameters, direction="write"
+                ),
+                "settings": convert_and_respect_annotation_metadata(
+                    object_=settings, annotation=typing.Optional[PromptSettings], direction="write"
+                ),
+                "blocks": convert_and_respect_annotation_metadata(
+                    object_=blocks, annotation=typing.Sequence[PromptBlock], direction="write"
+                ),
+                "functions": convert_and_respect_annotation_metadata(
+                    object_=functions,
+                    annotation=typing.Optional[typing.Sequence[FunctionDefinition]],
+                    direction="write",
+                ),
+                "expand_meta": convert_and_respect_annotation_metadata(
+                    object_=expand_meta, annotation=typing.Optional[AdHocExpandMeta], direction="write"
+                ),
+            },
+            headers={
+                "content-type": "application/json",
+            },
+            request_options=request_options,
+            omit=OMIT,
+        ) as _response:
+
+            async def stream() -> AsyncHttpResponse[typing.AsyncIterator[AdHocExecutePromptEvent]]:
+                try:
+                    if 200 <= _response.status_code < 300:
+
+                        async def _iter():
+                            async for _text in _response.aiter_lines():
+                                try:
+                                    if len(_text) == 0:
+                                        continue
+                                    yield typing.cast(
+                                        AdHocExecutePromptEvent,
+                                        parse_obj_as(
+                                            type_=AdHocExecutePromptEvent,  # type: ignore
+                                            object_=json.loads(_text),
+                                        ),
+                                    )
+                                except Exception:
+                                    pass
+                            return
+
+                        return AsyncHttpResponse(response=_response, data=_iter())
+                    await _response.aread()
+                    if _response.status_code == 400:
+                        raise BadRequestError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    if _response.status_code == 403:
+                        raise ForbiddenError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    if _response.status_code == 500:
+                        raise InternalServerError(
+                            typing.cast(
+                                typing.Optional[typing.Any],
+                                parse_obj_as(
+                                    type_=typing.Optional[typing.Any],  # type: ignore
+                                    object_=_response.json(),
+                                ),
+                            )
+                        )
+                    _response_json = _response.json()
+                except JSONDecodeError:
+                    raise ApiError(status_code=_response.status_code, body=_response.text)
+                raise ApiError(status_code=_response.status_code, body=_response_json)
+
+            yield await stream()

vellum/client/resources/container_images/client.py

@@ -33,7 +33,7 @@ class ContainerImagesClient:
         """
         return self._raw_client
 
-    def list_(
+    def list(
         self,
         *,
         limit: typing.Optional[int] = None,
@@ -73,12 +73,7 @@ class ContainerImagesClient:
         )
         client.container_images.list()
         """
-        response = self._raw_client.list(
-            limit=limit,
-            offset=offset,
-            ordering=ordering,
-            request_options=request_options,
-        )
+        response = self._raw_client.list(limit=limit, offset=offset, ordering=ordering, request_options=request_options)
         return response.data
 
     def retrieve(self, id: str, *, request_options: typing.Optional[RequestOptions] = None) -> ContainerImageRead:
@@ -110,10 +105,7 @@ class ContainerImagesClient:
             id="id",
         )
         """
-        response = self._raw_client.retrieve(
-            id,
-            request_options=request_options,
-        )
+        response = self._raw_client.retrieve(id, request_options=request_options)
         return response.data
 
     def docker_service_token(self, *, request_options: typing.Optional[RequestOptions] = None) -> DockerServiceToken:
@@ -138,9 +130,7 @@ class ContainerImagesClient:
         )
         client.container_images.docker_service_token()
         """
-        response = self._raw_client.docker_service_token(
-            request_options=request_options,
-        )
+        response = self._raw_client.docker_service_token(request_options=request_options)
         return response.data
 
     def push_container_image(
@@ -182,12 +172,7 @@ class ContainerImagesClient:
             tags=["tags", "tags"],
         )
         """
-        response = self._raw_client.push_container_image(
-            name=name,
-            sha=sha,
-            tags=tags,
-            request_options=request_options,
-        )
+        response = self._raw_client.push_container_image(name=name, sha=sha, tags=tags, request_options=request_options)
         return response.data
 
 
@@ -206,7 +191,7 @@ class AsyncContainerImagesClient:
         """
         return self._raw_client
 
-    async def list_(
+    async def list(
         self,
         *,
         limit: typing.Optional[int] = None,
@@ -255,10 +240,7 @@ class AsyncContainerImagesClient:
         asyncio.run(main())
         """
         response = await self._raw_client.list(
-            limit=limit,
-            offset=offset,
-            ordering=ordering,
-            request_options=request_options,
+            limit=limit, offset=offset, ordering=ordering, request_options=request_options
         )
         return response.data
 
@@ -299,10 +281,7 @@ class AsyncContainerImagesClient:
 
         asyncio.run(main())
         """
-        response = await self._raw_client.retrieve(
-            id,
-            request_options=request_options,
-        )
+        response = await self._raw_client.retrieve(id, request_options=request_options)
         return response.data
 
     async def docker_service_token(
@@ -337,9 +316,7 @@ class AsyncContainerImagesClient:
 
         asyncio.run(main())
         """
-        response = await self._raw_client.docker_service_token(
-            request_options=request_options,
-        )
+        response = await self._raw_client.docker_service_token(request_options=request_options)
         return response.data
 
     async def push_container_image(
@@ -390,9 +367,6 @@ class AsyncContainerImagesClient:
         asyncio.run(main())
         """
         response = await self._raw_client.push_container_image(
-            name=name,
-            sha=sha,
-            tags=tags,
-            request_options=request_options,
+            name=name, sha=sha, tags=tags, request_options=request_options
         )
         return response.data