Source code for taimoe.platform.resources.observability

"""Observability resource for uploading runtime spans."""

from __future__ import annotations

from typing import TYPE_CHECKING

from taimoe.platform.types.observability import SpanBatch, SpanEvent

if TYPE_CHECKING:
    from taimoe.platform.client.client import TaimoeClient




class ObservabilityResource:
    """Resource for interacting with /api/v1/observability endpoints."""

    def __init__(self, client: TaimoeClient) -> None:
        self._client = client



    def submit_spans(self, batch: SpanBatch) -> None:
        """Synchronously upload a batch of runtime spans."""
        self._client._request(
            "POST",
            "/api/v1/observability/spans/batch",
            json=batch.model_dump(mode="json"),
        )




    async def submit_spans_async(self, batch: SpanBatch) -> None:
        """Asynchronously upload a batch of runtime spans."""
        await self._client._request_async(
            "POST",
            "/api/v1/observability/spans/batch",
            json=batch.model_dump(mode="json"),
        )




    def submit_span(self, span: SpanEvent) -> None:
        """Upload one span as a single-span batch."""
        self.submit_spans(_single_span_batch(span))




    async def submit_span_async(self, span: SpanEvent) -> None:
        """Asynchronously upload one span as a single-span batch."""
        await self.submit_spans_async(_single_span_batch(span))




def _single_span_batch(span: SpanEvent) -> SpanBatch:
    return SpanBatch(
        runtime_id=span.runtime_id,
        trace_id=span.trace_id,
        session_id=span.session_id,
        spans=(span,),
    )