vibe-client 0.1.0__py3-none-any.whl

vibe_client/client.py

@@ -0,0 +1,498 @@
+# Code generated by smithy-python-codegen DO NOT EDIT.
+
+import asyncio
+from asyncio import Future, iscoroutine, sleep
+from copy import copy, deepcopy
+from dataclasses import replace
+import logging
+from typing import Any, Awaitable, Callable, cast
+
+from smithy_core import URI
+from smithy_core.deserializers import DeserializeableShape
+from smithy_core.endpoints import EndpointResolverParams
+from smithy_core.exceptions import SmithyRetryException
+from smithy_core.interceptors import (
+    InputContext,
+    Interceptor,
+    InterceptorChain,
+    OutputContext,
+    RequestContext,
+    ResponseContext,
+)
+from smithy_core.interfaces.exceptions import HasFault
+from smithy_core.interfaces.retries import RetryErrorInfo, RetryErrorType
+from smithy_core.schemas import APIOperation
+from smithy_core.serializers import SerializeableShape
+from smithy_core.types import TypedProperties
+from smithy_http.aio.interfaces import HTTPRequest, HTTPResponse
+from smithy_http.interfaces import HTTPRequestConfiguration
+from smithy_http.plugins import user_agent_plugin
+
+from .config import Config, Plugin
+from .deserialize import _deserialize_query_agent
+from .models import QUERY_AGENT, QueryAgentInput, QueryAgentOutput, ServiceError
+from .serialize import _serialize_query_agent
+from smithy_http import Field
+
+
+
+logger = logging.getLogger(__name__)
+
+class Vibe:
+    """
+    Client for Vibe
+
+    :param config: Optional configuration for the client. Here you can set things like the
+        endpoint for HTTP services or auth credentials.
+
+    :param plugins: A list of callables that modify the configuration dynamically. These
+        can be used to set defaults, for example.
+    """
+    def __init__(self, config: Config | None = None, plugins: list[Plugin] | None = None):
+        self._config = config or Config()
+
+        client_plugins: list[Plugin] = [
+            user_agent_plugin,
+        ]
+        if plugins:
+            client_plugins.extend(plugins)
+
+        for plugin in client_plugins:
+            plugin(self._config)
+
+    async def query_agent(self, input: QueryAgentInput, plugins: list[Plugin] | None = None) -> QueryAgentOutput:
+        """
+        Invokes the QueryAgent operation.
+
+        :param input: The operation's input.
+
+        :param plugins: A list of callables that modify the configuration dynamically.
+            Changes made by these plugins only apply for the duration of the operation
+            execution and will not affect any other operation invocations.
+
+        """
+        operation_plugins: list[Plugin] = [
+
+        ]
+        if plugins:
+            operation_plugins.extend(plugins)
+
+        return await self._execute_operation(
+            input=input,
+            plugins=operation_plugins,
+            serialize=_serialize_query_agent,
+            deserialize=_deserialize_query_agent,
+            config=self._config,
+            operation=QUERY_AGENT,
+        )
+
+    def _classify_error(
+        self,
+        *,
+        error: Exception,
+        context: ResponseContext[Any, HTTPRequest, HTTPResponse | None]
+    ) -> RetryErrorInfo:
+        logger.debug("Classifying error: %s", error)
+
+        if not isinstance(error, HasFault) and not context.transport_response:
+            return RetryErrorInfo(error_type=RetryErrorType.TRANSIENT)
+
+        if context.transport_response:
+            if context.transport_response.status in [429, 503]:
+                retry_after = None
+                retry_header = context.transport_response.fields["retry-after"]
+                if retry_header and retry_header.values:
+                    retry_after = float(retry_header.values[0])
+                return RetryErrorInfo(error_type=RetryErrorType.THROTTLING, retry_after_hint=retry_after)
+
+            if context.transport_response.status >= 500:
+                return RetryErrorInfo(error_type=RetryErrorType.SERVER_ERROR)
+
+        error_type = RetryErrorType.CLIENT_ERROR
+        if isinstance(error, HasFault) and error.fault == "server":
+            error_type = RetryErrorType.SERVER_ERROR
+
+        return RetryErrorInfo(error_type=error_type)
+
+    async def _execute_operation[Input: SerializeableShape, Output: DeserializeableShape](
+        self,
+        input: Input,
+        plugins: list[Plugin],
+        serialize: Callable[[Input, Config], Awaitable[HTTPRequest]],
+        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
+        config: Config,
+        operation: APIOperation[Input, Output],
+        request_future: Future[RequestContext[Any, HTTPRequest]] | None = None,
+        response_future: Future[HTTPResponse] | None = None,
+    ) -> Output:
+        try:
+            return await self._handle_execution(
+                input, plugins, serialize, deserialize, config, operation,
+                request_future, response_future,
+            )
+        except Exception as e:
+            if request_future is not None and not request_future.done():
+                request_future.set_exception(ServiceError(e))
+            if response_future is not None and not response_future.done():
+                response_future.set_exception(ServiceError(e))
+
+            # Make sure every exception that we throw is an instance of ServiceError so
+            # customers can reliably catch everything we throw.
+            if not isinstance(e, ServiceError):
+                raise ServiceError(e) from e
+            raise
+
+    async def _handle_execution[Input: SerializeableShape, Output: DeserializeableShape](
+        self,
+        input: Input,
+        plugins: list[Plugin],
+        serialize: Callable[[Input, Config], Awaitable[HTTPRequest]],
+        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
+        config: Config,
+        operation: APIOperation[Input, Output],
+        request_future: Future[RequestContext[Any, HTTPRequest]] | None,
+        response_future: Future[HTTPResponse] | None,
+    ) -> Output:
+        operation_name = operation.schema.id.name
+        logger.debug('Making request for operation "%s" with parameters: %s', operation_name, input)
+        config = deepcopy(config)
+        for plugin in plugins:
+            plugin(config)
+
+        input_context = InputContext(request=input, properties=TypedProperties({"config": config}))
+        transport_request: HTTPRequest | None = None
+        output_context: OutputContext[Input, Output, HTTPRequest | None, HTTPResponse | None] | None = None
+
+        client_interceptors = cast(
+            list[Interceptor[Input, Output, HTTPRequest, HTTPResponse]], list(config.interceptors)
+        )
+        interceptor_chain = InterceptorChain(client_interceptors)
+
+        try:
+            # Step 1: Invoke read_before_execution
+            interceptor_chain.read_before_execution(input_context)
+
+            # Step 2: Invoke the modify_before_serialization hooks
+            input_context = replace(
+                input_context,
+                request=interceptor_chain.modify_before_serialization(input_context)
+            )
+
+            # Step 3: Invoke the read_before_serialization hooks
+            interceptor_chain.read_before_serialization(input_context)
+
+            # Step 4: Serialize the request
+            logger.debug("Serializing request for: %s", input_context.request)
+            transport_request = await serialize(input_context.request, config)
+            request_context = RequestContext(
+                request=input_context.request,
+                transport_request=transport_request,
+                properties=input_context.properties,
+            )
+            logger.debug("Serialization complete. Transport request: %s", request_context.transport_request)
+
+            # Step 5: Invoke read_after_serialization
+            interceptor_chain.read_after_serialization(request_context)
+
+            # Step 6: Invoke modify_before_retry_loop
+            request_context = replace(
+                request_context,
+                transport_request=interceptor_chain.modify_before_retry_loop(request_context)
+            )
+
+            # Step 7: Acquire the retry token.
+            retry_strategy = config.retry_strategy
+            retry_token = retry_strategy.acquire_initial_retry_token()
+
+            while True:
+                # Make an attempt
+                output_context = await self._handle_attempt(
+                    deserialize,
+                    interceptor_chain,
+                    replace(
+                      request_context,
+                      transport_request = copy(request_context.transport_request)
+                    ),
+                    config,
+                    operation,
+                    request_future,
+                )
+
+                if isinstance(output_context.response, Exception):
+                    # Step 7u: Reacquire retry token if the attempt failed
+                    try:
+                        retry_token = retry_strategy.refresh_retry_token_for_retry(
+                            token_to_renew=retry_token,
+                            error_info=self._classify_error(
+                                error=output_context.response,
+                                context=output_context,
+                            )
+                        )
+                    except SmithyRetryException:
+                        raise output_context.response
+                    logger.debug(
+                        "Retry needed. Attempting request #%s in %.4f seconds.",
+                        retry_token.retry_count + 1,
+                        retry_token.retry_delay
+                    )
+                    await sleep(retry_token.retry_delay)
+                    current_body = output_context.transport_request.body
+                    if (seek := getattr(current_body, "seek", None)) is not None:
+                        if iscoroutine((result := seek(0))):
+                            await result
+                else:
+                    # Step 8: Invoke record_success
+                    retry_strategy.record_success(token=retry_token)
+                    if response_future is not None:
+                        response_future.set_result(
+                            output_context.transport_response  # type: ignore
+                        )
+                    break
+        except Exception as e:
+            if output_context is not None:
+                logger.exception("Exception occurred while handling: %s", output_context.response)
+                output_context = replace(output_context, response=e)
+            else:
+                output_context = OutputContext(
+                    request=input_context.request,
+                    response=e,
+                    transport_request=transport_request,
+                    transport_response=None,
+                    properties=input_context.properties
+                )
+
+        return await self._finalize_execution(interceptor_chain, output_context)
+
+    async def _handle_attempt[Input: SerializeableShape, Output: DeserializeableShape](
+        self,
+        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
+        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
+        context: RequestContext[Input, HTTPRequest],
+        config: Config,
+        operation: APIOperation[Input, Output],
+        request_future: Future[RequestContext[Input, HTTPRequest]] | None,
+    ) -> OutputContext[Input, Output, HTTPRequest, HTTPResponse | None]:
+        transport_response: HTTPResponse | None = None
+        try:
+            # Step 7a: Invoke read_before_attempt
+            interceptor.read_before_attempt(context)
+
+            # Step 7f: Invoke endpoint_resolver.resolve_endpoint
+            endpoint_resolver_parameters = EndpointResolverParams(
+                operation=operation,
+                input=context.request,
+                context=context.properties
+            )
+            logger.debug("Calling endpoint resolver with parameters: %s", endpoint_resolver_parameters)
+            endpoint = await config.endpoint_resolver.resolve_endpoint(
+                endpoint_resolver_parameters
+            )
+            logger.debug("Endpoint resolver result: %s", endpoint)
+            if not endpoint.uri.path:
+                path = ""
+            elif endpoint.uri.path.endswith("/"):
+                path = endpoint.uri.path[:-1]
+            else:
+                path = endpoint.uri.path
+            if context.transport_request.destination.path:
+                path += context.transport_request.destination.path
+            context.transport_request.destination = URI(
+                scheme=endpoint.uri.scheme,
+                host=context.transport_request.destination.host + endpoint.uri.host,
+                path=path,
+                port=endpoint.uri.port,
+                query=context.transport_request.destination.query,
+            )
+
+            if (headers := endpoint.properties.get("headers")) is not None:
+                context.transport_request.fields.extend(headers)
+
+            # Step 7g: Invoke modify_before_signing
+            context = replace(
+                context,
+                transport_request=interceptor.modify_before_signing(context)
+            )
+
+            # Step 7h: Invoke read_before_signing
+            interceptor.read_before_signing(context)
+
+            # Step 7j: Invoke read_after_signing
+            interceptor.read_after_signing(context)
+
+            # Step 7k: Invoke modify_before_transmit
+            context = replace(
+                context,
+                transport_request=interceptor.modify_before_transmit(context)
+            )
+
+            # Step 7l: Invoke read_before_transmit
+            interceptor.read_before_transmit(context)
+
+            # Step 7m: Invoke http_client.send
+            request_config = config.http_request_config or HTTPRequestConfiguration()
+            logger.debug("HTTP request config: %s", request_config)
+            logger.debug("Sending HTTP request: %s", context.transport_request)
+
+            if request_future is not None:
+                response_task = asyncio.create_task(config.http_client.send(
+                    request=context.transport_request,
+                    request_config=request_config,
+                ))
+                request_future.set_result(context)
+                transport_response = await response_task
+            else:
+                transport_response = await config.http_client.send(
+                    request=context.transport_request,
+                    request_config=request_config,
+                )
+
+            response_context = ResponseContext(
+                request=context.request,
+                transport_request=context.transport_request,
+                transport_response=transport_response,
+                properties=context.properties
+            )
+            logger.debug("Received HTTP response: %s", response_context.transport_response)
+
+            # Step 7n: Invoke read_after_transmit
+            interceptor.read_after_transmit(response_context)
+
+            # Step 7o: Invoke modify_before_deserialization
+            response_context = replace(
+                response_context,
+                transport_response=interceptor.modify_before_deserialization(response_context)
+            )
+
+            # Step 7p: Invoke read_before_deserialization
+            interceptor.read_before_deserialization(response_context)
+
+            # Step 7q: deserialize
+            logger.debug("Deserializing transport response: %s", response_context.transport_response)
+            output = await deserialize(
+                response_context.transport_response, config
+            )
+            output_context = OutputContext(
+                request=response_context.request,
+                response=output,
+                transport_request=response_context.transport_request,
+                transport_response=response_context.transport_response,
+                properties=response_context.properties
+            )
+            logger.debug("Deserialization complete. Response: %s", output_context.response)
+
+            # Step 7r: Invoke read_after_deserialization
+            interceptor.read_after_deserialization(output_context)
+        except Exception as e:
+            output_context: OutputContext[Input, Output, HTTPRequest, HTTPResponse] = OutputContext(
+                request=context.request,
+                response=e,  # type: ignore
+                transport_request=context.transport_request,
+                transport_response=transport_response,
+                properties=context.properties
+            )
+
+        return await self._finalize_attempt(interceptor, output_context)
+
+    async def _finalize_attempt[Input: SerializeableShape, Output: DeserializeableShape](
+        self,
+        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
+        context: OutputContext[Input, Output, HTTPRequest, HTTPResponse | None],
+    ) -> OutputContext[Input, Output, HTTPRequest, HTTPResponse | None]:
+        # Step 7s: Invoke modify_before_attempt_completion
+        try:
+            context = replace(
+                context,
+                response=interceptor.modify_before_attempt_completion(context)
+            )
+        except Exception as e:
+            logger.exception("Exception occurred while handling: %s", context.response)
+            context = replace(context, response=e)
+
+        # Step 7t: Invoke read_after_attempt
+        try:
+            interceptor.read_after_attempt(context)
+        except Exception as e:
+            context = replace(context, response=e)
+
+        return context
+
+    async def _finalize_execution[Input: SerializeableShape, Output: DeserializeableShape](
+        self,
+        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
+        context: OutputContext[Input, Output, HTTPRequest | None, HTTPResponse | None],
+    ) -> Output:
+        try:
+            # Step 9: Invoke modify_before_completion
+            context = replace(
+                context,
+                response=interceptor.modify_before_completion(context)
+            )
+
+            # Step 10: Invoke trace_probe.dispatch_events
+            try:
+                pass
+            except Exception as e:
+                # log and ignore exceptions
+                logger.exception("Exception occurred while dispatching trace events: %s", e)
+                pass
+        except Exception as e:
+            logger.exception("Exception occurred while handling: %s", context.response)
+            context = replace(context, response=e)
+
+        # Step 11: Invoke read_after_execution
+        try:
+            interceptor.read_after_execution(context)
+        except Exception as e:
+            context = replace(context, response=e)
+
+        # Step 12: Return / throw
+        if isinstance(context.response, Exception):
+            raise context.response
+
+        # We may want to add some aspects of this context to the output types so we can
+        # return it to the end-users.
+        return context.response
+
+
+
+class ApiKeyInterceptor(Interceptor):
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+
+    def modify_before_signing(self, context: RequestContext) -> HTTPRequest:
+        """Add the x-api-key header to the request before it's signed."""
+
+        request = context.transport_request
+        # Add the x-api-key header directly
+        request.fields.set_field(Field(name="x-api-key", values=[str(self.api_key)]))
+        return context.transport_request
+
+class VibeClient(Vibe):
+    """
+    An extension of the Vibe client with built-in API key authentication.
+    """
+
+    def __init__(self, api_key: str, endpoint: str = "https://us-east-1.vibe.api.astartech.ai/", **kwargs):
+        """
+        Initialize a Vibe client with API key.
+
+        Args:
+            api_key: The API key for authentication
+            endpoint: The API endpoint URI (default: https://us-east-1.vibe.api.astartech.ai/)
+            **kwargs: Additional arguments to pass to the Vibe constructor
+
+        Raises:
+            ValueError: If api_key is None or empty
+        """
+        # Check that api_key is not None or empty
+        if not api_key:
+            raise ValueError("api_key cannot be None or empty")
+
+        # Create config with API key interceptor
+        config = Config(
+            endpoint_uri=endpoint,
+            interceptors=[ApiKeyInterceptor(api_key)]
+        )
+
+        # Initialize the parent Vibe class with our config
+        super().__init__(config=config, **kwargs)

vibe_client/config.py

@@ -0,0 +1,74 @@
+# Code generated by smithy-python-codegen DO NOT EDIT.
+
+from dataclasses import dataclass
+from typing import Any, Callable, TypeAlias, Union
+
+from smithy_core.aio.endpoints import StaticEndpointResolver
+from smithy_core.aio.interfaces import EndpointResolver as _EndpointResolver
+from smithy_core.interceptors import Interceptor
+from smithy_core.interfaces import URI
+from smithy_core.interfaces.retries import RetryStrategy
+from smithy_core.retries import SimpleRetryStrategy
+from smithy_http.aio.aiohttp import AIOHTTPClient
+from smithy_http.aio.interfaces import HTTPClient
+from smithy_http.interfaces import HTTPRequestConfiguration
+
+from .models import QueryAgentInput, QueryAgentOutput
+
+
+_ServiceInterceptor = Union[Interceptor[QueryAgentInput, QueryAgentOutput, Any, Any]]
+@dataclass(init=False)
+class Config:
+    """Configuration for Vibe."""
+
+    endpoint_resolver: _EndpointResolver
+    endpoint_uri: str | URI | None
+    http_client: HTTPClient
+    http_request_config: HTTPRequestConfiguration | None
+    interceptors: list[_ServiceInterceptor]
+    retry_strategy: RetryStrategy
+
+    def __init__(
+        self,
+        *,
+        endpoint_resolver: _EndpointResolver | None = None,
+        endpoint_uri: str | URI | None = None,
+        http_client: HTTPClient | None = None,
+        http_request_config: HTTPRequestConfiguration | None = None,
+        interceptors: list[_ServiceInterceptor] | None = None,
+        retry_strategy: RetryStrategy | None = None,
+    ):
+        """Constructor.
+
+        :param endpoint_resolver:
+             The endpoint resolver used to resolve the final endpoint per-operation based on
+             the configuration.
+
+        :param endpoint_uri:
+             A static URI to route requests to.
+
+        :param http_client:
+             The HTTP client used to make requests.
+
+        :param http_request_config:
+             Configuration for individual HTTP requests.
+
+        :param interceptors:
+             The list of interceptors, which are hooks that are called during the execution
+             of a request.
+
+        :param retry_strategy:
+             The retry strategy for issuing retry tokens and computing retry delays.
+
+        """
+        self.endpoint_resolver = endpoint_resolver or StaticEndpointResolver()
+        self.endpoint_uri = endpoint_uri
+        self.http_client = http_client or AIOHTTPClient()
+        self.http_request_config = http_request_config
+        self.interceptors = interceptors or []
+        self.retry_strategy = retry_strategy or SimpleRetryStrategy()
+
+#
+# A callable that allows customizing the config object on each request.
+#
+Plugin: TypeAlias = Callable[[Config], None]

vibe_client/deserialize.py

@@ -0,0 +1,115 @@
+# Code generated by smithy-python-codegen DO NOT EDIT.
+
+import json
+from typing import Any
+
+from smithy_core.documents import DocumentValue
+from smithy_core.types import TimestampFormat
+from smithy_http.aio.interfaces import HTTPResponse
+from smithy_http.aio.restjson import parse_rest_json_error_info
+from smithy_json import JSONCodec
+
+from .config import Config
+from .models import (
+    ApiError,
+    QueryAgentOutput,
+    UnauthorizedException,
+    UnknownApiError,
+    ValidationException,
+    VibeValidationException,
+)
+
+
+async def _deserialize_query_agent(http_response: HTTPResponse, config: Config) -> QueryAgentOutput:
+    if http_response.status != 200 and http_response.status >= 300:
+        raise await _deserialize_error_query_agent(http_response, config)
+
+    kwargs: dict[str, Any] = {}
+
+    body = await http_response.consume_body_async()
+    if body:
+        codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
+        deserializer = codec.create_deserializer(body)
+        body_kwargs = QueryAgentOutput.deserialize_kwargs(deserializer)
+        kwargs.update(body_kwargs)
+
+    return QueryAgentOutput(**kwargs)
+
+async def _deserialize_error_query_agent(http_response: HTTPResponse, config: Config) -> ApiError:
+    code, message, parsed_body = await parse_rest_json_error_info(http_response)
+
+    match code.lower():
+        case "unauthorizedexception":
+            return await _deserialize_error_unauthorized_exception(http_response, config, parsed_body, message)
+
+        case "validationexception":
+            return await _deserialize_error_validation_exception(http_response, config, parsed_body, message)
+
+        case "vibevalidationexception":
+            return await _deserialize_error_vibe_validation_exception(http_response, config, parsed_body, message)
+
+        case _:
+            return UnknownApiError(f"{code}: {message}")
+
+async def _deserialize_error_validation_exception(
+    http_response: HTTPResponse,
+    config: Config,
+    parsed_body: dict[str, DocumentValue] | None,
+    default_message: str,
+) -> ValidationException:
+    kwargs: dict[str, Any] = {"message": default_message}
+
+    if parsed_body is None:
+        body = await http_response.consume_body_async()
+    else:
+        body = json.dumps(parsed_body).encode('utf-8')
+
+    if body:
+        codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
+        deserializer = codec.create_deserializer(body)
+        body_kwargs = ValidationException.deserialize_kwargs(deserializer)
+        kwargs.update(body_kwargs)
+
+    return ValidationException(**kwargs)
+
+async def _deserialize_error_unauthorized_exception(
+    http_response: HTTPResponse,
+    config: Config,
+    parsed_body: dict[str, DocumentValue] | None,
+    default_message: str,
+) -> UnauthorizedException:
+    kwargs: dict[str, Any] = {"message": default_message}
+
+    if parsed_body is None:
+        body = await http_response.consume_body_async()
+    else:
+        body = json.dumps(parsed_body).encode('utf-8')
+
+    if body:
+        codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
+        deserializer = codec.create_deserializer(body)
+        body_kwargs = UnauthorizedException.deserialize_kwargs(deserializer)
+        kwargs.update(body_kwargs)
+
+    return UnauthorizedException(**kwargs)
+
+async def _deserialize_error_vibe_validation_exception(
+    http_response: HTTPResponse,
+    config: Config,
+    parsed_body: dict[str, DocumentValue] | None,
+    default_message: str,
+) -> VibeValidationException:
+    kwargs: dict[str, Any] = {"message": default_message}
+
+    if parsed_body is None:
+        body = await http_response.consume_body_async()
+    else:
+        body = json.dumps(parsed_body).encode('utf-8')
+
+    if body:
+        codec = JSONCodec(default_timestamp_format=TimestampFormat.EPOCH_SECONDS)
+        deserializer = codec.create_deserializer(body)
+        body_kwargs = VibeValidationException.deserialize_kwargs(deserializer)
+        kwargs.update(body_kwargs)
+
+    return VibeValidationException(**kwargs)