scale-gp-beta 0.1.0a23__py3-none-any.whl → 0.1.0a25__py3-none-any.whl

scale_gp_beta/__init__.py

@@ -37,7 +37,7 @@ from ._exceptions import (
     UnprocessableEntityError,
     APIResponseValidationError,
 )
-from ._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient
+from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
 
 __all__ = [
@@ -80,6 +80,7 @@ __all__ = [
     "DEFAULT_CONNECTION_LIMITS",
     "DefaultHttpxClient",
     "DefaultAsyncHttpxClient",
+    "DefaultAioHttpClient",
 ]
 
 if not _t.TYPE_CHECKING:

scale_gp_beta/_base_client.py

@@ -1289,6 +1289,24 @@ class _DefaultAsyncHttpxClient(httpx.AsyncClient):
         super().__init__(**kwargs)
 
 
+try:
+    import httpx_aiohttp
+except ImportError:
+
+    class _DefaultAioHttpClient(httpx.AsyncClient):
+        def __init__(self, **_kwargs: Any) -> None:
+            raise RuntimeError("To use the aiohttp client you must have installed the package with the `aiohttp` extra")
+else:
+
+    class _DefaultAioHttpClient(httpx_aiohttp.HttpxAiohttpClient):  # type: ignore
+        def __init__(self, **kwargs: Any) -> None:
+            kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
+            kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
+            kwargs.setdefault("follow_redirects", True)
+
+            super().__init__(**kwargs)
+
+
 if TYPE_CHECKING:
     DefaultAsyncHttpxClient = httpx.AsyncClient
     """An alias to `httpx.AsyncClient` that provides the same defaults that this SDK
@@ -1297,8 +1315,12 @@ if TYPE_CHECKING:
     This is useful because overriding the `http_client` with your own instance of
     `httpx.AsyncClient` will result in httpx's defaults being used, not ours.
     """
+
+    DefaultAioHttpClient = httpx.AsyncClient
+    """An alias to `httpx.AsyncClient` that changes the default HTTP transport to `aiohttp`."""
 else:
     DefaultAsyncHttpxClient = _DefaultAsyncHttpxClient
+    DefaultAioHttpClient = _DefaultAioHttpClient
 
 
 class AsyncHttpxClientWrapper(DefaultAsyncHttpxClient):

scale_gp_beta/_client.py

@@ -21,7 +21,17 @@ from ._types import (
 )
 from ._utils import is_given, get_async_library
 from ._version import __version__
-from .resources import spans, models, datasets, inference, completions, evaluations, dataset_items, evaluation_items
+from .resources import (
+    spans,
+    models,
+    datasets,
+    inference,
+    questions,
+    completions,
+    evaluations,
+    dataset_items,
+    evaluation_items,
+)
 from ._streaming import Stream as Stream, AsyncStream as AsyncStream
 from ._exceptions import APIStatusError, SGPClientError
 from ._base_client import (
@@ -54,6 +64,7 @@ class SGPClient(SyncAPIClient):
     completions: completions.CompletionsResource
     chat: chat.ChatResource
     inference: inference.InferenceResource
+    questions: questions.QuestionsResource
     files: files.FilesResource
     models: models.ModelsResource
     datasets: datasets.DatasetsResource
@@ -157,6 +168,7 @@ class SGPClient(SyncAPIClient):
         self.completions = completions.CompletionsResource(self)
         self.chat = chat.ChatResource(self)
         self.inference = inference.InferenceResource(self)
+        self.questions = questions.QuestionsResource(self)
         self.files = files.FilesResource(self)
         self.models = models.ModelsResource(self)
         self.datasets = datasets.DatasetsResource(self)
@@ -281,6 +293,7 @@ class AsyncSGPClient(AsyncAPIClient):
     completions: completions.AsyncCompletionsResource
     chat: chat.AsyncChatResource
     inference: inference.AsyncInferenceResource
+    questions: questions.AsyncQuestionsResource
     files: files.AsyncFilesResource
     models: models.AsyncModelsResource
     datasets: datasets.AsyncDatasetsResource
@@ -384,6 +397,7 @@ class AsyncSGPClient(AsyncAPIClient):
         self.completions = completions.AsyncCompletionsResource(self)
         self.chat = chat.AsyncChatResource(self)
         self.inference = inference.AsyncInferenceResource(self)
+        self.questions = questions.AsyncQuestionsResource(self)
         self.files = files.AsyncFilesResource(self)
         self.models = models.AsyncModelsResource(self)
         self.datasets = datasets.AsyncDatasetsResource(self)
@@ -509,6 +523,7 @@ class SGPClientWithRawResponse:
         self.completions = completions.CompletionsResourceWithRawResponse(client.completions)
         self.chat = chat.ChatResourceWithRawResponse(client.chat)
         self.inference = inference.InferenceResourceWithRawResponse(client.inference)
+        self.questions = questions.QuestionsResourceWithRawResponse(client.questions)
         self.files = files.FilesResourceWithRawResponse(client.files)
         self.models = models.ModelsResourceWithRawResponse(client.models)
         self.datasets = datasets.DatasetsResourceWithRawResponse(client.datasets)
@@ -523,6 +538,7 @@ class AsyncSGPClientWithRawResponse:
         self.completions = completions.AsyncCompletionsResourceWithRawResponse(client.completions)
         self.chat = chat.AsyncChatResourceWithRawResponse(client.chat)
         self.inference = inference.AsyncInferenceResourceWithRawResponse(client.inference)
+        self.questions = questions.AsyncQuestionsResourceWithRawResponse(client.questions)
         self.files = files.AsyncFilesResourceWithRawResponse(client.files)
         self.models = models.AsyncModelsResourceWithRawResponse(client.models)
         self.datasets = datasets.AsyncDatasetsResourceWithRawResponse(client.datasets)
@@ -537,6 +553,7 @@ class SGPClientWithStreamedResponse:
         self.completions = completions.CompletionsResourceWithStreamingResponse(client.completions)
         self.chat = chat.ChatResourceWithStreamingResponse(client.chat)
         self.inference = inference.InferenceResourceWithStreamingResponse(client.inference)
+        self.questions = questions.QuestionsResourceWithStreamingResponse(client.questions)
         self.files = files.FilesResourceWithStreamingResponse(client.files)
         self.models = models.ModelsResourceWithStreamingResponse(client.models)
         self.datasets = datasets.DatasetsResourceWithStreamingResponse(client.datasets)
@@ -551,6 +568,7 @@ class AsyncSGPClientWithStreamedResponse:
         self.completions = completions.AsyncCompletionsResourceWithStreamingResponse(client.completions)
         self.chat = chat.AsyncChatResourceWithStreamingResponse(client.chat)
         self.inference = inference.AsyncInferenceResourceWithStreamingResponse(client.inference)
+        self.questions = questions.AsyncQuestionsResourceWithStreamingResponse(client.questions)
         self.files = files.AsyncFilesResourceWithStreamingResponse(client.files)
         self.models = models.AsyncModelsResourceWithStreamingResponse(client.models)
         self.datasets = datasets.AsyncDatasetsResourceWithStreamingResponse(client.datasets)

scale_gp_beta/_version.py

@@ -1,4 +1,4 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 __title__ = "scale_gp_beta"
-__version__ = "0.1.0-alpha.23"  # x-release-please-version
+__version__ = "0.1.0-alpha.25"  # x-release-please-version

scale_gp_beta/lib/tracing/span.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import logging
+from copy import deepcopy
 from typing import TYPE_CHECKING, Type, Optional
+from threading import RLock
 from typing_extensions import override
 
 from scale_gp_beta.types.span_upsert_batch_params import Item as SpanCreateRequest
@@ -64,8 +66,8 @@ class BaseSpan:
         self._group_id = group_id
         self._span_id: str = span_id or generate_span_id()
         self._parent_span_id = parent_span_id
-        self.start_time: Optional[str] = None
-        self.end_time: Optional[str] = None
+        self._start_time: Optional[str] = None
+        self._end_time: Optional[str] = None
         self._input: SpanInputParam = input or {}
         self._output: SpanOutputParam = output or {}
         self._metadata: SpanMetadataParam = metadata or {}
@@ -74,6 +76,7 @@ class BaseSpan:
         self._queue_manager = queue_manager
 
         self._contextvar_token: Optional[contextvars.Token[Optional[BaseSpan]]] = None
+        self._lock = RLock()
 
     def start(self) -> None:
         pass
@@ -113,13 +116,33 @@ class BaseSpan:
         return self._span_type
 
     # with setters
+    @property
+    def start_time(self) -> Optional[str]:
+        return self._start_time
+
+    @start_time.setter
+    def start_time(self, value: Optional[str]) -> None:
+        with self._lock:
+            self._start_time = value
+
+    @property
+    def end_time(self) -> Optional[str]:
+        return self._end_time
+
+    @end_time.setter
+    def end_time(self, value: Optional[str]) -> None:
+        with self._lock:
+            self._end_time = value
+
     @property
     def metadata(self) -> SpanMetadataParam:
         return self._metadata
 
     @metadata.setter
     def metadata(self, value: SpanMetadataParam) -> None:
-        self._metadata = value
+        # this does not protect against span.metadata["foo"] = "bar" which uses the getter, ditto input and output
+        with self._lock:
+            self._metadata = value
 
     @property
     def input(self) -> SpanInputParam:
@@ -127,7 +150,8 @@ class BaseSpan:
 
     @input.setter
     def input(self, value: SpanInputParam) -> None:
-        self._input = value
+        with self._lock:
+            self._input = value
 
     @property
     def output(self) -> SpanOutputParam:
@@ -135,7 +159,8 @@ class BaseSpan:
 
     @output.setter
     def output(self, value: SpanOutputParam) -> None:
-        self._output = value
+        with self._lock:
+            self._output = value
 
     def set_error(
             self,
@@ -144,12 +169,13 @@ class BaseSpan:
             exception: Optional[BaseException] = None,
     ) -> None:
         # Naively record details in metadata for now, note that error capture only supported in context manager
-        exception_type = type(exception).__name__ if exception else None
-        exception_message = str(exception) if exception else None
-        self._status = "ERROR"
-        self.metadata["error"] = True
-        self.metadata["error_type"] = error_type or exception_type
-        self.metadata["error_message"] = error_message or exception_message
+        with self._lock:
+            exception_type = type(exception).__name__ if exception else None
+            exception_message = str(exception) if exception else None
+            self._status = "ERROR"
+            self.metadata["error"] = True
+            self.metadata["error_type"] = error_type or exception_type
+            self.metadata["error_message"] = error_message or exception_message
 
     def __enter__(self) -> BaseSpan:
         self.start()
@@ -167,32 +193,35 @@ class BaseSpan:
         self.end()
 
     def to_request_params(self) -> SpanCreateRequest:
-        if self.start_time is None:
-            raise ParamsCreationError("No start time specified")
-
-        request_data = SpanCreateRequest(
-            name=self.name,
-            id=self.span_id,
-            trace_id=self.trace_id,
-            start_timestamp=self.start_time,
-            input=self.input,
-            output=self.output,
-            metadata=self.metadata,
-            status=self.status,
-            type=self.span_type
-        )
+        with self._lock:
+            if self.start_time is None:
+                raise ParamsCreationError("No start time specified")
+
+            request_data = SpanCreateRequest(
+                name=self.name,
+                id=self.span_id,
+                trace_id=self.trace_id,
+                start_timestamp=self.start_time,
+                input=self.input,
+                output=self.output,
+                metadata=self.metadata,
+                status=self.status,
+                type=self.span_type
+            )
 
-        if self.end_time is not None:
-            request_data["end_timestamp"] = self.end_time
+            if self.end_time is not None:
+                request_data["end_timestamp"] = self.end_time
 
-        # parent_span_id is optional (root spans)
-        if self.parent_span_id is not None:
-            request_data["parent_id"] = self.parent_span_id
+            # parent_span_id is optional (root spans)
+            if self.parent_span_id is not None:
+                request_data["parent_id"] = self.parent_span_id
 
-        if self.group_id is not None:
-            request_data["group_id"] = self.group_id
+            if self.group_id is not None:
+                request_data["group_id"] = self.group_id
 
-        return request_data
+            # ensure no future changes to metadata, input or output changes request_data, full isolation
+            request_data = deepcopy(request_data)
+            return request_data
 
     @override
     def __repr__(self) -> str:
@@ -288,13 +317,14 @@ class Span(BaseSpan):
         Sets the `start_time`, reports the span start to the `TraceQueueManager`
         , and registers this span as the current span.
         """
-        if self.start_time is not None:
-            log.warning(f"Span {self.name}: {self.span_id} has already started at {self.start_time}")
-            return
+        with self._lock:
+            if self.start_time is not None:
+                log.warning(f"Span {self.name}: {self.span_id} has already started at {self.start_time}")
+                return
 
-        self.start_time = iso_timestamp()
-        self._queue_manager.report_span_start(self)
-        self._contextvar_token = Scope.set_current_span(self)
+            self.start_time = iso_timestamp()
+            self._queue_manager.report_span_start(self)
+            self._contextvar_token = Scope.set_current_span(self)
 
     @override
     def end(self) -> None:
@@ -304,19 +334,20 @@ class Span(BaseSpan):
         `TraceQueueManager` for queuing and export, and resets this span from the
         `Scope`.
         """
-        if self.end_time is not None:
-            log.warning(f"Span {self.name}: {self.span_id} has already ended at {self.end_time}")
-            return
-        if self._contextvar_token is None:
-            log.warning(
-                (
-                    f"Span {self.name}: {self.span_id} attempting to end without a valid context token. "
-                    "Was start() called and completed successfully?"
+        with self._lock:
+            if self.end_time is not None:
+                log.warning(f"Span {self.name}: {self.span_id} has already ended at {self.end_time}")
+                return
+            if self._contextvar_token is None:
+                log.warning(
+                    (
+                        f"Span {self.name}: {self.span_id} attempting to end without a valid context token. "
+                        "Was start() called and completed successfully?"
+                    )
                 )
-            )
-            return
+                return
 
-        self.end_time = iso_timestamp()
-        self._queue_manager.report_span_end(self)
-        Scope.reset_current_span(self._contextvar_token)
-        self._contextvar_token = None
+            self.end_time = iso_timestamp()
+            self._queue_manager.report_span_end(self)
+            Scope.reset_current_span(self._contextvar_token)
+            self._contextvar_token = None

scale_gp_beta/lib/tracing/trace_queue_manager.py

@@ -30,7 +30,7 @@ class TraceQueueManager:
     """Manage trace and spans queue
     Store spans in-memory until the threshold has been reached then flush to server.
 
-    Optionally provide a client, if unprovided we will attempt to create a default client.
+    Optionally provide a client, if unprovided, we will attempt to create a default client.
     """
 
     def __init__(
@@ -41,6 +41,7 @@ class TraceQueueManager:
         trigger_queue_size: int = DEFAULT_TRIGGER_QUEUE_SIZE,
         trigger_cadence: float = DEFAULT_TRIGGER_CADENCE,
         retries: int = DEFAULT_RETRIES,
+        worker_enabled: Optional[bool] = None,
     ):
         self._client = client
         self._attempted_local_client_creation = False
@@ -54,7 +55,9 @@ class TraceQueueManager:
         self._shutdown_event = threading.Event()
         self._queue: queue.Queue[Span] = queue.Queue(maxsize=max_queue_size)
 
-        if not is_disabled():
+        self._worker_enabled = worker_enabled if worker_enabled is not None else not is_disabled()
+
+        if self._worker_enabled:
             self._worker = threading.Thread(daemon=True, target=self._run)
             self._worker.start()
 
@@ -66,7 +69,7 @@ class TraceQueueManager:
         self._client = client
 
     def shutdown(self, timeout: Optional[float] = None) -> None:
-        if is_disabled():
+        if not self._worker_enabled:
             log.debug("No worker to shutdown")
             return
         log.info(f"Shutting down trace queue manager, joining worker thread with timeout {timeout}")
@@ -92,6 +95,7 @@ class TraceQueueManager:
 
     def enqueue(self, span: "Span") -> None:
         try:
+            # Should this be a deep copy of span instead? Currently is a reference
             self._queue.put_nowait(span)
         except queue.Full:
             log.warning(f"Queue full, ignoring span {span.span_id}")

scale_gp_beta/lib/tracing/tracing.py

@@ -6,7 +6,7 @@ from .util import is_disabled
 from .scope import Scope
 from .trace import Trace, BaseTrace, NoOpTrace
 from .types import SpanInputParam, SpanOutputParam, SpanTypeLiterals, SpanMetadataParam
-from .trace_queue_manager import tracing_queue_manager
+from .trace_queue_manager import TraceQueueManager, tracing_queue_manager
 
 log: logging.Logger = logging.getLogger(__name__)
 
@@ -58,6 +58,7 @@ def create_trace(
         span_id: Optional[str] = None,
         trace_id: Optional[str] = None,
         group_id: Optional[str] = None,
+        queue_manager: Optional[TraceQueueManager] = None,
 ) -> BaseTrace:
     """Creates a new trace and root span instance.
 
@@ -83,6 +84,8 @@ def create_trace(
                                   If None, a unique trace ID will be generated.
                                   Max length is 38 characters.
         group_id (Optional[str]): An optional, id to group traces.
+        queue_manager (Optional[TraceQueueManager], optional): An optional `TraceQueueManager`.
+            Useful for when you need explicit control of flushing and client behavior.
 
     Returns:
         BaseTrace: A `Trace` instance if tracing is enabled, or a `NoOpTrace`
@@ -109,7 +112,7 @@ def create_trace(
     if active_trace is not None:
         log.warning(f"Trace with id {active_trace.trace_id} is already active. Creating a new trace anyways.")
 
-    queue_manager = tracing_queue_manager()
+    queue_manager = queue_manager or tracing_queue_manager()
     trace = Trace(
         name=name,
         trace_id=trace_id,
@@ -136,6 +139,7 @@ def create_span(
     parent_id: Optional[str] = None,
     trace_id: Optional[str] = None,
     group_id: Optional[str] = None,
+    queue_manager: Optional[TraceQueueManager] = None,
 ) -> BaseSpan:
     """Creates a new span instance.
 
@@ -171,6 +175,8 @@ def create_span(
         trace_id (Optional[str], optional): A `Trace` id. Used for explicit control.
             Default to trace id fetched from the active scope.
         group_id (Optional[str]): An optional, id to group traces.
+        queue_manager (Optional[TraceQueueManager], optional): An optional `TraceQueueManager`.
+            Useful for when you need explicit control of flushing and client behavior.
 
     Returns:
         BaseSpan: A `Span` instance if tracing is enabled and a valid trace context
@@ -213,7 +219,7 @@ def create_span(
         log.debug(f"Attempting to create a span with no trace")
         return noop_span
 
-    queue_manager = tracing_queue_manager()
+    queue_manager = queue_manager or tracing_queue_manager()
     span = Span(
         name=name,
         span_id=span_id,

scale_gp_beta/lib/tracing/types.py

@@ -13,7 +13,7 @@ SpanInputParam = Dict[str, Any]
 SpanOutputParam = Dict[str, Any]
 SpanMetadataParam = Dict[str, Any]
 
-SpanStatusLiterals = Literal["SUCCESS", "ERROR"]
+SpanStatusLiterals = Literal["SUCCESS", "ERROR", "CANCELED"]
 
 SpanTypeLiterals = Literal[
         "TEXT_INPUT",
@@ -30,8 +30,6 @@ SpanTypeLiterals = Literal[
         "DOCUMENT_SEARCH",
         "DOCUMENT_PROMPT",
         "CUSTOM",
-        "INPUT_GUARDRAIL",
-        "OUTPUT_GUARDRAIL",
         "CODE_EXECUTION",
         "DATA_MANIPULATION",
         "EVALUATION",

scale_gp_beta/resources/__init__.py

@@ -48,6 +48,14 @@ from .inference import (
     InferenceResourceWithStreamingResponse,
     AsyncInferenceResourceWithStreamingResponse,
 )
+from .questions import (
+    QuestionsResource,
+    AsyncQuestionsResource,
+    QuestionsResourceWithRawResponse,
+    AsyncQuestionsResourceWithRawResponse,
+    QuestionsResourceWithStreamingResponse,
+    AsyncQuestionsResourceWithStreamingResponse,
+)
 from .completions import (
     CompletionsResource,
     AsyncCompletionsResource,
@@ -100,6 +108,12 @@ __all__ = [
     "AsyncInferenceResourceWithRawResponse",
     "InferenceResourceWithStreamingResponse",
     "AsyncInferenceResourceWithStreamingResponse",
+    "QuestionsResource",
+    "AsyncQuestionsResource",
+    "QuestionsResourceWithRawResponse",
+    "AsyncQuestionsResourceWithRawResponse",
+    "QuestionsResourceWithStreamingResponse",
+    "AsyncQuestionsResourceWithStreamingResponse",
     "FilesResource",
     "AsyncFilesResource",
     "FilesResourceWithRawResponse",

scale_gp_beta/resources/chat/completions.py

@@ -536,6 +536,7 @@ class CompletionsResource(SyncAPIResource):
             "model_zoo",
             "bedrock",
             "xai",
+            "fireworks_ai",
         ]
         | NotGiven = NOT_GIVEN,
         sort_order: Literal["asc", "desc"] | NotGiven = NOT_GIVEN,
@@ -548,7 +549,7 @@ class CompletionsResource(SyncAPIResource):
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> CompletionModelsResponse:
         """
-        Chat Completions
+        List Chat Completion Models
 
         Args:
           extra_headers: Send extra headers
@@ -1090,6 +1091,7 @@ class AsyncCompletionsResource(AsyncAPIResource):
             "model_zoo",
             "bedrock",
             "xai",
+            "fireworks_ai",
         ]
         | NotGiven = NOT_GIVEN,
         sort_order: Literal["asc", "desc"] | NotGiven = NOT_GIVEN,
@@ -1102,7 +1104,7 @@ class AsyncCompletionsResource(AsyncAPIResource):
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> CompletionModelsResponse:
         """
-        Chat Completions
+        List Chat Completion Models
 
         Args:
           extra_headers: Send extra headers

scale_gp_beta/resources/models.py

@@ -317,6 +317,7 @@ class ModelsResource(SyncAPIResource):
             "model_zoo",
             "bedrock",
             "xai",
+            "fireworks_ai",
         ]
         | NotGiven = NOT_GIVEN,
         name: str | NotGiven = NOT_GIVEN,
@@ -689,6 +690,7 @@ class AsyncModelsResource(AsyncAPIResource):
             "model_zoo",
             "bedrock",
             "xai",
+            "fireworks_ai",
         ]
         | NotGiven = NOT_GIVEN,
         name: str | NotGiven = NOT_GIVEN,