from http import HTTPStatus
import chariot._apis
from chariot.config import getLogger
from chariot.tracker import _utils, models
__all__ = [
"create_tracker",
"get_tracker",
"get_trackers",
"update_tracker",
"delete_tracker",
]
log = getLogger(__name__)
[docs]
def create_tracker(
request_body: models.NewCreateTrackerRequest,
) -> models.Tracker | None:
"""Create a new tracker.
:param request_body: The create tracker specification.
:type request_body: NewCreateTrackerRequest
:return: The tracker.
:rtype: Optional[models.Tracker]
"""
response = chariot._apis.models.tracker_api.tracker_post(body=request_body.model_dump())
if not response:
return None
return _utils.convert_to_dataclass(response.model_dump(), models.Tracker)
[docs]
def get_tracker(tracker_id: str) -> models.Tracker | None:
"""Get an existing tracker.
:param tracker_id: The tracker id.
:type tracker_id: str
:return: The tracker.
:rtype: Optional[models.Tracker]
"""
response = chariot._apis.models.tracker_api.tracker_id_get(id=tracker_id)
if not response:
return None
return _utils.convert_to_dataclass(response.model_dump(), models.Tracker)
[docs]
def get_trackers(
project_id: str,
name: str | None = None,
fuzzy_name: str | None = None,
sort: str = "created_at:desc",
offset: int = 0,
limit: int = 10,
) -> list[models.Tracker]:
"""Get existing trackers.
:param project_id: The project id
:type project_id: str
:param name: The exact name of the tracker.
:type name: str
:param fuzzy_name: Search for trackers given a fuzzy name.
:type fuzzy_name: str
:param sort: The column and direction to sort by (defaults to created_at:desc).
:type sort: str
:param limit: Limit the number of trackers returned (defaults to 10).
:type limit: int
:param offset: Offset the starting position of trackers to be returned (defaults to 0).
:type offset: int
:return: The trackers.
:rtype: list[models.Tracker]
"""
response = chariot._apis.models.tracker_api.tracker_get(
project_ids=project_id,
tracker_name=name,
fuzzy_name=fuzzy_name,
sort=sort,
offset=offset,
limit=limit,
)
if not response:
return []
return [_utils.convert_to_dataclass(d.model_dump(), models.Tracker) for d in response.data]
[docs]
def update_tracker(
request_body: models.NewUpdateTrackerRequest,
) -> models.NewStoreTracksRequest | None:
"""Update an existing tracker and optionally wait for the track output.
If `wait_for_output` is False, this function always returns None. Otherwise,
a successful update (HTTP 200) returns a `NewStoreTracksRequest`. To ensure
synchronous behavior within the tracking service, the API handler for updates
must wait for the track output message to be published and received.
The time required to process an update and receive the output can vary.
If the `wait_time` expires before the output is received, a REQUEST_TIMEOUT (408)
status is returned. This does not indicate that the tracker update failed, only
that the output could not be delivered within the `wait_time` window.
The tracks should always be available within the inference-store.
:param request_body: The update tracker specification.
:type request_body: NewUpdateTrackerRequest
:return: The track outputs.
:rtype: Optional[models.NewStoreTracksRequest]
"""
response = chariot._apis.trackproducer.tracker_api.tracker_update_post_with_http_info(
body=request_body.model_dump()
)
if not response or not request_body.wait_for_output:
return None
match response.status_code:
case HTTPStatus.OK:
return _utils.convert_to_dataclass(
response.data.data.model_dump(), models.NewStoreTracksRequest
)
case HTTPStatus.REQUEST_TIMEOUT:
log.error("Timed out waiting for tracker update response")
return None
case _:
raise RuntimeError(f"Unexpected status code {response.status_code}: {response.data}")
[docs]
def delete_tracker(tracker_id: str):
"""Delete an existing tracker.
:param tracker_id: The tracker id.
:type tracker_id: str
"""
return chariot._apis.models.tracker_api.tracker_id_delete(id=tracker_id)