import json
import numpy as np
import pandas as pd

import pytz
import datetime

import requests
from requests.auth import HTTPBasicAuth

import wallaroo

wl = wallaroo.Client()

# this will also format the timezone in the parsing section
timezone = "US/Central"

selected_timezone = pytz.timezone(timezone)

# Define the start and end times of 10:00 to 10:15
data_start = selected_timezone.localize(datetime.datetime(2025, 7, 14, 10, 0, 0))
data_end = selected_timezone.localize(datetime.datetime(2025, 7, 14, 10, 15, 00))

# this is the URL to get prometheus metrics
query_url = f"{wl.api_endpoint}/v1/metrics/api/v1/query_range"

import time
# Retrieve the token 
headers = wl.auth.auth_header()

# Convert to UTC and get the Unix timestamps
start_timestamp = int(data_start.astimezone(pytz.UTC).timestamp())
end_timestamp = int(data_end.astimezone(pytz.UTC).timestamp())    

pipeline_name = "llama-3-1-8b-pipeline" # the name of the pipeline
deploy_id = 210 # the deployment id
step = "5m" # the step of the calculation

query_ttft = f'histogram_quantile(0.99, sum(rate(vllm:time_to_first_token_seconds_bucket{{kubernetes_namespace="{pipeline_name}-{deploy_id}"}}[{step}])) by (le)) * 1000'
print(query_ttft)

#request parameters
params_ttft = {
    'query': query_ttft,
    'start': start_timestamp,
    'end': end_timestamp,
    'step': step
}

response_rps = requests.get(query_url, headers=headers, params=params_ttft)

if response_rps.status_code == 200:
    #print("Requests Per Second Data:")
    result = response_rps.json()
    print(result)
else:
    print("Failed to fetch TTFT data:", response_rps.status_code, response_rps.text)

histogram_quantile(0.99, sum(rate(vllm:time_to_first_token_seconds_bucket{kubernetes_namespace="llama-3-1-8b-pipeline-210"}[5m])) by (le)) * 1000
{'status': 'success', 'data': {'resultType': 'matrix', 'result': [{'metric': {}, 'values': [[1752505500, '48.45656000000012'], [1752505800, '39.800000000000004'], [1752506100, 'NaN']]}]}}

# this will also format the timezone in the parsing section
timezone = "US/Central"

selected_timezone = pytz.timezone(timezone)

# Define the start and end times of 10:00 to 10:15
data_start = selected_timezone.localize(datetime.datetime(2025, 7, 14, 10, 0, 0))
data_end = selected_timezone.localize(datetime.datetime(2025, 7, 14, 10, 15, 00))

# this is the URL to get prometheus metrics
query_url = f"{wl.api_endpoint}/v1/metrics/api/v1/query_range"

import time
# Retrieve the token 
headers = wl.auth.auth_header()

# Convert to UTC and get the Unix timestamps
start_timestamp = int(data_start.astimezone(pytz.UTC).timestamp())
end_timestamp = int(data_end.astimezone(pytz.UTC).timestamp())    

pipeline_name = "llama-3-1-8b-pipeline" # the name of the pipeline
deploy_id = 210 # the deployment id
step = "5m" # the step of the calculation

query_tps = f'sum by (kubernetes_namespace) (rate(vllm:generation_tokens_total{{kubernetes_namespace="{pipeline_name}-{deploy_id}"}}[{step}]))'
print(query_tps)

#request parameters
params_ttft = {
    'query': query_tps,
    'start': start_timestamp,
    'end': end_timestamp,
    'step': step
}

response_rps = requests.get(query_url, headers=headers, params=params_ttft)

if response_rps.status_code == 200:
    #print("Requests Per Second Data:")
    result = response_rps.json()
    print(result)
else:
    print("Failed to fetch TTFT data:", response_rps.status_code, response_rps.text)

sum by (kubernetes_namespace) (rate(vllm:generation_tokens_total{kubernetes_namespace="llama-3-1-8b-pipeline-210"}[5m]))
{'status': 'success', 'data': {'resultType': 'matrix', 'result': [{'metric': {'kubernetes_namespace': 'llama-3-1-8b-pipeline-210'}, 'values': [[1752505200, '0'], [1752505500, '0.6707186440677967'], [1752505800, '0.6779661016949152'], [1752506100, '0']]}]}}