forked from Zipstack/unstract
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathviews.py
More file actions
63 lines (49 loc) · 2.58 KB
/
Copy pathviews.py
File metadata and controls
63 lines (49 loc) · 2.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import logging
from django.http import HttpRequest
from rest_framework import status, viewsets
from rest_framework.decorators import action
from rest_framework.exceptions import APIException, ValidationError
from rest_framework.response import Response
from .constants import UsageKeys
from .helper import UsageHelper
from .serializers import GetUsageSerializer
logger = logging.getLogger(__name__)
class UsageView(viewsets.ModelViewSet):
"""Viewset for managing Usage-related operations."""
@action(detail=True, methods=["get"])
def get_token_usage(self, request: HttpRequest) -> Response:
"""Retrieves the aggregated token usage for a given run_id.
This method validates the 'run_id' query parameter, aggregates the token
usage statistics for the specified run_id, and returns the results.
Args:
request (HttpRequest): The HTTP request object containing the
query parameters.
Returns:
Response: A Response object containing the aggregated token usage data
with HTTP 200 OK status if successful, or an error message and
appropriate HTTP status if an error occurs.
Raises:
ValidationError: If the 'run_id' query parameter is missing or invalid.
APIException: If an unexpected error occurs during the execution.
"""
try:
# Validate the query parameters using the serializer
# This ensures that 'run_id' is present and valid
serializer = GetUsageSerializer(data=self.request.query_params)
serializer.is_valid(raise_exception=True)
run_id = serializer.validated_data.get(UsageKeys.RUN_ID)
# Retrieve aggregated token count for the given run_id.
result: dict = UsageHelper.get_aggregated_token_count(run_id=run_id)
# Log the successful completion of the operation
logger.info(f"Token usage retrieved successfully for run_id: {run_id}")
# Return the result
return Response(status=status.HTTP_200_OK, data=result)
except ValidationError as e:
# Handle validation errors specifically
logger.error(f"Validation error: {e.detail}")
raise ValidationError(detail=f"Validation error: {str(e)}")
except Exception as e:
# Handle any other exceptions that might occur during the execution
error_msg = "An unexpected error occurred while fetching the token usage"
logger.error(f"{error_msg}: {e}")
raise APIException(detail=error_msg)