General: Track token usage of LLM service requests

src/main/java/de/tum/cit/aet/artemis/core/domain/LLMRequest.java

@@ -0,0 +1,4 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+public record LLMRequest(String model, int numInputTokens, float costPerMillionInputToken, int numOutputTokens, float costPerMillionOutputToken, String pipelineId) {
+}

src/main/java/de/tum/cit/aet/artemis/core/domain/LLMServiceType.java

@@ -0,0 +1,8 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+/**
+ * Enum representing different types of LLM (Large Language Model) services used in the system.
+ */
+public enum LLMServiceType {
+    IRIS, ATHENA
+}

src/main/java/de/tum/cit/aet/artemis/core/domain/LLMTokenUsageRequest.java

@@ -0,0 +1,95 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.ManyToOne;
+import jakarta.persistence.Table;
+
+import org.hibernate.annotations.Cache;
+import org.hibernate.annotations.CacheConcurrencyStrategy;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+
+@Entity
+@Table(name = "llm_token_usage_request")
+@Cache(usage = CacheConcurrencyStrategy.NONSTRICT_READ_WRITE)
+@JsonInclude(JsonInclude.Include.NON_EMPTY)
+public class LLMTokenUsageRequest extends DomainObject {
+
+    @Column(name = "model")
+    private String model;
+
+    @Column(name = "service_pipeline_id")
+    private String servicePipelineId;
+
+    @Column(name = "num_input_tokens")
+    private int numInputTokens;
+
+    @Column(name = "cost_per_million_input_tokens")
+    private float costPerMillionInputTokens;
+
+    @Column(name = "num_output_tokens")
+    private int numOutputTokens;
+
+    @Column(name = "cost_per_million_output_tokens")
+    private float costPerMillionOutputTokens;
+
+    @ManyToOne
+    private LLMTokenUsageTrace trace;
+
+    public String getModel() {
+        return model;
+    }
+
+    public void setModel(String model) {
+        this.model = model;
+    }
+
+    public String getServicePipelineId() {
+        return servicePipelineId;
+    }
+
+    public void setServicePipelineId(String servicePipelineId) {
+        this.servicePipelineId = servicePipelineId;
+    }
+
+    public float getCostPerMillionInputTokens() {
+        return costPerMillionInputTokens;
+    }
+
+    public void setCostPerMillionInputTokens(float costPerMillionInputToken) {
+        this.costPerMillionInputTokens = costPerMillionInputToken;
+    }
+
+    public float getCostPerMillionOutputTokens() {
+        return costPerMillionOutputTokens;
+    }
+
+    public void setCostPerMillionOutputTokens(float costPerMillionOutputToken) {
+        this.costPerMillionOutputTokens = costPerMillionOutputToken;
+    }
+
+    public int getNumInputTokens() {
+        return numInputTokens;
+    }
+
+    public void setNumInputTokens(int numInputTokens) {
+        this.numInputTokens = numInputTokens;
+    }
+
+    public int getNumOutputTokens() {
+        return numOutputTokens;
+    }
+
+    public void setNumOutputTokens(int numOutputTokens) {
+        this.numOutputTokens = numOutputTokens;
+    }
+
+    public LLMTokenUsageTrace getTrace() {
+        return trace;
+    }
+
+    public void setTrace(LLMTokenUsageTrace trace) {
+        this.trace = trace;
+    }
+}

src/main/java/de/tum/cit/aet/artemis/core/domain/LLMTokenUsageTrace.java

@@ -0,0 +1,108 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+import java.time.ZonedDateTime;
+import java.util.HashSet;
+import java.util.Set;
+
+import jakarta.annotation.Nullable;
+import jakarta.persistence.CascadeType;
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.EnumType;
+import jakarta.persistence.Enumerated;
+import jakarta.persistence.FetchType;
+import jakarta.persistence.OneToMany;
+import jakarta.persistence.Table;
+
+import org.hibernate.annotations.Cache;
+import org.hibernate.annotations.CacheConcurrencyStrategy;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+
+@Entity
+@Table(name = "llm_token_usage_trace")
+@Cache(usage = CacheConcurrencyStrategy.NONSTRICT_READ_WRITE)
+@JsonInclude(JsonInclude.Include.NON_EMPTY)
+public class LLMTokenUsageTrace extends DomainObject {
+
+    @Column(name = "service")
+    @Enumerated(EnumType.STRING)
+    private LLMServiceType serviceType;
+
+    @Nullable
+    @Column(name = "course_id")
+    private Long courseId;
+
+    @Nullable
+    @Column(name = "exercise_id")
+    private Long exerciseId;
+
+    @Column(name = "user_id")
+    private Long userId;
+
+    @Column(name = "time")
+    private ZonedDateTime time = ZonedDateTime.now();
+
+    @Nullable
+    @Column(name = "iris_message_id")
+    private Long irisMessageId;
+
+    @OneToMany(mappedBy = "trace", fetch = FetchType.LAZY, cascade = CascadeType.ALL, orphanRemoval = true)
+    private Set<LLMTokenUsageRequest> llmRequests = new HashSet<>();
+
+    public LLMServiceType getServiceType() {
+        return serviceType;
+    }
+
+    public void setServiceType(LLMServiceType serviceType) {
+        this.serviceType = serviceType;
+    }
+
+    public Long getCourseId() {
+        return courseId;
+    }
+
+    public void setCourseId(Long courseId) {
+        this.courseId = courseId;
+    }
+
+    public Long getExerciseId() {
+        return exerciseId;
+    }
+
+    public void setExerciseId(Long exerciseId) {
+        this.exerciseId = exerciseId;
+    }
+
+    public long getUserId() {
+        return userId;
+    }
+
+    public void setUserId(long userId) {
+        this.userId = userId;
+    }
+
+    public ZonedDateTime getTime() {
+        return time;
+    }
+
+    public void setTime(ZonedDateTime time) {
+        this.time = time;
+    }
+
+    public Set<LLMTokenUsageRequest> getLLMRequests() {
+        return llmRequests;
+    }
+
+    public void setLlmRequests(Set<LLMTokenUsageRequest> llmRequests) {
+        this.llmRequests = llmRequests;
+    }
+
+    public Long getIrisMessageId() {
+        return irisMessageId;
+    }
+
+    public void setIrisMessageId(Long messageId) {
+        this.irisMessageId = messageId;
+    }
+}

src/main/java/de/tum/cit/aet/artemis/core/repository/LLMTokenUsageRequestRepository.java

@@ -0,0 +1,14 @@
+package de.tum.cit.aet.artemis.core.repository;
+
+import static de.tum.cit.aet.artemis.core.config.Constants.PROFILE_CORE;
+
+import org.springframework.context.annotation.Profile;
+import org.springframework.stereotype.Repository;
+
+import de.tum.cit.aet.artemis.core.domain.LLMTokenUsageRequest;
+import de.tum.cit.aet.artemis.core.repository.base.ArtemisJpaRepository;
+
+@Profile(PROFILE_CORE)
+@Repository
+public interface LLMTokenUsageRequestRepository extends ArtemisJpaRepository<LLMTokenUsageRequest, Long> {
+}

src/main/java/de/tum/cit/aet/artemis/core/repository/LLMTokenUsageTraceRepository.java

@@ -0,0 +1,14 @@
+package de.tum.cit.aet.artemis.core.repository;
+
+import static de.tum.cit.aet.artemis.core.config.Constants.PROFILE_CORE;
+
+import org.springframework.context.annotation.Profile;
+import org.springframework.stereotype.Repository;
+
+import de.tum.cit.aet.artemis.core.domain.LLMTokenUsageTrace;
+import de.tum.cit.aet.artemis.core.repository.base.ArtemisJpaRepository;
+
+@Profile(PROFILE_CORE)
+@Repository
+public interface LLMTokenUsageTraceRepository extends ArtemisJpaRepository<LLMTokenUsageTrace, Long> {
+}

src/main/java/de/tum/cit/aet/artemis/core/service/LLMTokenUsageService.java

@@ -0,0 +1,133 @@
+package de.tum.cit.aet.artemis.core.service;
+
+import static de.tum.cit.aet.artemis.core.config.Constants.PROFILE_CORE;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import org.springframework.context.annotation.Profile;
+import org.springframework.stereotype.Service;
+
+import de.tum.cit.aet.artemis.core.domain.LLMRequest;
+import de.tum.cit.aet.artemis.core.domain.LLMServiceType;
+import de.tum.cit.aet.artemis.core.domain.LLMTokenUsageRequest;
+import de.tum.cit.aet.artemis.core.domain.LLMTokenUsageTrace;
+import de.tum.cit.aet.artemis.core.repository.LLMTokenUsageRequestRepository;
+import de.tum.cit.aet.artemis.core.repository.LLMTokenUsageTraceRepository;
+
+/**
+ * Service for managing the LLMTokenUsage by all LLMs in Artemis
+ */
+@Profile(PROFILE_CORE)
+@Service
+public class LLMTokenUsageService {
+
+    private final LLMTokenUsageTraceRepository llmTokenUsageTraceRepository;
+
+    private final LLMTokenUsageRequestRepository llmTokenUsageRequestRepository;
+
+    public LLMTokenUsageService(LLMTokenUsageTraceRepository llmTokenUsageTraceRepository, LLMTokenUsageRequestRepository llmTokenUsageRequestRepository) {
+        this.llmTokenUsageTraceRepository = llmTokenUsageTraceRepository;
+        this.llmTokenUsageRequestRepository = llmTokenUsageRequestRepository;
+    }
+
+    /**
+     * Saves the token usage to the database.
+     * This method records the usage of tokens by various LLM services in the system.
+     *
+     * @param llmRequests     List of LLM requests containing details about the token usage.
+     * @param serviceType     Type of the LLM service (e.g., IRIS, GPT-3).
+     * @param builderFunction A function that takes an LLMTokenUsageBuilder and returns a modified LLMTokenUsageBuilder.
+     *                            This function is used to set additional properties on the LLMTokenUsageTrace object, such as
+     *                            the course ID, user ID, exercise ID, and Iris message ID.
+     *                            Example usage:
+     *                            builder -> builder.withCourse(courseId).withUser(userId)
+     * @return The saved LLMTokenUsageTrace object, which includes the details of the token usage.
+     */
+    // TODO: this should ideally be done Async
+    public LLMTokenUsageTrace saveLLMTokenUsage(List<LLMRequest> llmRequests, LLMServiceType serviceType, Function<LLMTokenUsageBuilder, LLMTokenUsageBuilder> builderFunction) {
+        LLMTokenUsageTrace llmTokenUsageTrace = new LLMTokenUsageTrace();
+        llmTokenUsageTrace.setServiceType(serviceType);
+
+        LLMTokenUsageBuilder builder = builderFunction.apply(new LLMTokenUsageBuilder());
+        builder.getIrisMessageID().ifPresent(llmTokenUsageTrace::setIrisMessageId);
+        builder.getCourseID().ifPresent(llmTokenUsageTrace::setCourseId);
+        builder.getExerciseID().ifPresent(llmTokenUsageTrace::setExerciseId);
+        builder.getUserID().ifPresent(llmTokenUsageTrace::setUserId);
+
+        llmTokenUsageTrace.setLlmRequests(llmRequests.stream().map(LLMTokenUsageService::convertLLMRequestToLLMTokenUsageRequest)
+                .peek(llmTokenUsageRequest -> llmTokenUsageRequest.setTrace(llmTokenUsageTrace)).collect(Collectors.toSet()));
+
+        return llmTokenUsageTraceRepository.save(llmTokenUsageTrace);
+    }
+
+    private static LLMTokenUsageRequest convertLLMRequestToLLMTokenUsageRequest(LLMRequest llmRequest) {
+        LLMTokenUsageRequest llmTokenUsageRequest = new LLMTokenUsageRequest();
+        llmTokenUsageRequest.setModel(llmRequest.model());
+        llmTokenUsageRequest.setNumInputTokens(llmRequest.numInputTokens());
+        llmTokenUsageRequest.setNumOutputTokens(llmRequest.numOutputTokens());
+        llmTokenUsageRequest.setCostPerMillionInputTokens(llmRequest.costPerMillionInputToken());
+        llmTokenUsageRequest.setCostPerMillionOutputTokens(llmRequest.costPerMillionOutputToken());
+        llmTokenUsageRequest.setServicePipelineId(llmRequest.pipelineId());
+        return llmTokenUsageRequest;
+    }
+
+    // TODO: this should ideally be done Async
+    public void appendRequestsToTrace(List<LLMRequest> requests, LLMTokenUsageTrace trace) {
+        var requestSet = requests.stream().map(LLMTokenUsageService::convertLLMRequestToLLMTokenUsageRequest).peek(llmTokenUsageRequest -> llmTokenUsageRequest.setTrace(trace))
+                .collect(Collectors.toSet());
+        llmTokenUsageRequestRepository.saveAll(requestSet);
+    }
+
+    /**
+     * Class LLMTokenUsageBuilder to be used for saveLLMTokenUsage()
+     */
+    public static class LLMTokenUsageBuilder {
+
+        private Optional<Long> courseID = Optional.empty();
+
+        private Optional<Long> irisMessageID = Optional.empty();
+
+        private Optional<Long> exerciseID = Optional.empty();
+
+        private Optional<Long> userID = Optional.empty();
+
+        public LLMTokenUsageBuilder withCourse(Long courseID) {
+            this.courseID = Optional.ofNullable(courseID);
+            return this;
+        }
+
+        public LLMTokenUsageBuilder withIrisMessageID(Long irisMessageID) {
+            this.irisMessageID = Optional.ofNullable(irisMessageID);
+            return this;
+        }
+
+        public LLMTokenUsageBuilder withExercise(Long exerciseID) {
+            this.exerciseID = Optional.ofNullable(exerciseID);
+            return this;
+        }
+
+        public LLMTokenUsageBuilder withUser(Long userID) {
+            this.userID = Optional.ofNullable(userID);
+            return this;
+        }
+
+        public Optional<Long> getCourseID() {
+            return courseID;
+        }
+
+        public Optional<Long> getIrisMessageID() {
+            return irisMessageID;
+        }
+
+        public Optional<Long> getExerciseID() {
+            return exerciseID;
+        }
+
+        public Optional<Long> getUserID() {
+            return userID;
+        }
+    }
+}