Notes

package com.citi.cti.eap.puma.service.impl;

import com.citi.cti.eap.puma.config.OpenAiConfig;
import com.citi.cti.eap.puma.constants.OpenAiConstants;
import com.citi.cti.eap.puma.model.PumaApplicationProperties;
import com.citi.cti.eap.puma.model.openai.Cost;
import com.citi.cti.eap.puma.model.openai.OpenAICost;
import com.citi.cti.eap.puma.ms.client.PumaApplicationPropertiesClient;
import com.citi.cti.eap.puma.repository.OpenAiCostMetricsRepository;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;

import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneId;
import java.util.Date;
import java.util.List;

@Service
public class OpenAIMetricsServiceImpl {

private static final Logger logger = LoggerFactory.getLogger(OpenAIMetricsServiceImpl.class);
private final ObjectMapper objectMapper = new ObjectMapper();

@Autowired
private OpenAiConfig openAiConfig;

@Autowired
private PumaApplicationPropertiesClient papService;

@Autowired
private OpenAiCostMetricsRepository openAiCostMetricsRepository;

public Cost getCost(LocalDate start_time, LocalDate end_time) {
PumaApplicationProperties hostProperties = papService.getByPropertyNameAndCategoryType(
OpenAiConstants.OpenAi_COST_URL,
OpenAiConstants.PROPERTY_CATEGORY_TYPE_OpenAi);

PumaApplicationProperties bearerTokenProperties = papService.getByPropertyNameAndCategoryType(
OpenAiConstants.OpenAi_BEARER_TOKEN,
OpenAiConstants.PROPERTY_CATEGORY_TYPE_OpenAi);

ZoneId zoneId = ZoneId.of("America/New_York");
String url = hostProperties.getPropertyValue() + "?group_by=project_id,line_item";

if (start_time != null) {
url += "&start_time=" + start_time.atStartOfDay(zoneId).toEpochSecond();
}
if (end_time != null) {
url += "&end_time=" + end_time.atStartOfDay(zoneId).toEpochSecond();
}

logger.info("Calling OpenAI Cost API at URL: {}", url);

RestClient restClient = openAiConfig.restOpenAiClient();
Cost response = restClient.get()
.uri(url)
.accept(MediaType.APPLICATION_JSON)
.headers(httpHeaders -> {
httpHeaders.add("Authorization", "Bearer " + bearerTokenProperties.getPropertyValue());
httpHeaders.add("Content-Type", "application/json");
})
.retrieve().body(Cost.class);

try {
logger.info("Full cost response: {}", objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(response));
} catch (Exception e) {
logger.error("Error while logging cost response", e);
}

response.getData().forEach(data -> {
if (!data.getResults().isEmpty()) {
data.getResults().forEach(result -> {
OpenAICost cost = new OpenAICost();
cost.setStartTime(Instant.ofEpochSecond(data.getStartTime()).atZone(zoneId).toLocalDate());
cost.setEndTime(Instant.ofEpochSecond(data.getEndTime()).atZone(zoneId).toLocalDate());
cost.setType(result.getObject());
cost.setCurrency(result.getAmount().getCurrency());
cost.setAmount(result.getAmount().getValue());
cost.setOrgId(result.getOrganizationId());
cost.setProjectId(result.getProjectId());
cost.setLineItem(result.getLineItem());
cost.setCreatedDatetime(new Date());

List<OpenAICost> removed = openAiCostMetricsRepository
.deleteByStartTimeAndEndTimeAndOrgIdAndProjectId(
cost.getStartTime(),
cost.getEndTime(),
cost.getOrgId(),
cost.getProjectId()
);

logger.info("Deleted {} existing cost records", removed.size());
openAiCostMetricsRepository.save(cost);
});
}
});

return response;
}

public String getUsage(LocalDate start_time, LocalDate end_time) {
PumaApplicationProperties hostProperties = papService.getByPropertyNameAndCategoryType(
OpenAiConstants.OpenAi_USAGE_URL,
OpenAiConstants.PROPERTY_CATEGORY_TYPE_OpenAi);

PumaApplicationProperties bearerTokenProperties = papService.getByPropertyNameAndCategoryType(
OpenAiConstants.OpenAi_BEARER_TOKEN,
OpenAiConstants.PROPERTY_CATEGORY_TYPE_OpenAi);

ZoneId zoneId = ZoneId.of("America/New_York");
String url = hostProperties.getPropertyValue() + "?group_by=project_id,model";

if (start_time != null) {
url += "&start_time=" + start_time.atStartOfDay(zoneId).toEpochSecond();
}
if (end_time != null) {
url += "&end_time=" + end_time.atStartOfDay(zoneId).toEpochSecond();
}

logger.info("Calling OpenAI Usage API at URL: {}", url);

RestClient restClient = openAiConfig.restOpenAiClient();
String usageJson = restClient.get()
.uri(url)
.accept(MediaType.APPLICATION_JSON)
.headers(headers -> {
headers.add("Authorization", "Bearer " + bearerTokenProperties.getPropertyValue());
headers.add("Content-Type", "application/json");
})
.retrieve()
.body(String.class);

try {
logger.info("Full usage response: {}", objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(
objectMapper.readTree(usageJson)
));
} catch (Exception e) {
logger.error("Error while logging usage response", e);
}

return usageJson;
}
}