Iris: Track token usage of iris requests

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

@@ -0,0 +1,31 @@
+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 {
+    /** Athena service for preliminary feedback */
+    ATHENA_PRELIMINARY_FEEDBACK,
+    /** Athena service for feedback suggestions */
+    ATHENA_FEEDBACK_SUGGESTION,
+    /** Iris service for code feedback */
+    IRIS_CODE_FEEDBACK,
+    /** Iris service for course chat messages */
+    IRIS_CHAT_COURSE_MESSAGE,
+    /** Iris service for exercise chat messages */
+    IRIS_CHAT_EXERCISE_MESSAGE,
+    /** Iris service for interaction suggestions */
+    IRIS_INTERACTION_SUGGESTION,
+    /** Iris service for lecture chat messages */
+    IRIS_CHAT_LECTURE_MESSAGE,
+    /** Iris service for competency generation */
+    IRIS_COMPETENCY_GENERATION,
+    /** Iris service for citation pipeline */
+    IRIS_CITATION_PIPELINE,
+    /** Iris service for lecture retrieval pipeline */
+    IRIS_LECTURE_RETRIEVAL_PIPELINE,
+    /** Iris service for lecture ingestion */
+    IRIS_LECTURE_INGESTION,
+    /** Default value when the service type is not set */
+    NOT_SET
+}

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

@@ -0,0 +1,177 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+import java.time.ZonedDateTime;
+
+import jakarta.annotation.Nullable;
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.EnumType;
+import jakarta.persistence.Enumerated;
+import jakarta.persistence.JoinColumn;
+import jakarta.persistence.ManyToOne;
+import jakarta.persistence.Table;
+
+import org.hibernate.annotations.Cache;
+import org.hibernate.annotations.CacheConcurrencyStrategy;
+
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonInclude;
+
+import de.tum.cit.aet.artemis.exercise.domain.Exercise;
+import de.tum.cit.aet.artemis.iris.domain.message.IrisMessage;
+
+@Entity
+@Table(name = "llm_token_usage")
+@Cache(usage = CacheConcurrencyStrategy.NONSTRICT_READ_WRITE)
+@JsonInclude(JsonInclude.Include.NON_EMPTY)
+public class LLMTokenUsage extends DomainObject {
+
+    @Column(name = "service")
+    @Enumerated(EnumType.STRING)
+    private LLMServiceType serviceType;
+
+    @Column(name = "model")
+    private String model;
+
+    @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;
+
+    @Nullable
+    @ManyToOne
+    @JsonIgnore
+    @JoinColumn(name = "course_id")
+    private Course course;
+
+    @Nullable
+    @ManyToOne
+    @JsonIgnore
+    @JoinColumn(name = "exercise_id")
+    private Exercise exercise;
+
+    @Column(name = "user_id")
+    private long userId;
+
+    @Column(name = "time")
+    private ZonedDateTime time = ZonedDateTime.now();
+
+    @Column(name = "trace_id")
+    private String traceId;
+
+    @Nullable
+    @ManyToOne
+    @JsonIgnore
+    @JoinColumn(name = "iris_message_id")
+    private IrisMessage irisMessage;
+
+    public LLMServiceType getServiceType() {
+        return serviceType;
+    }
+
+    public void setServiceType(LLMServiceType serviceType) {
+        this.serviceType = serviceType;
+    }
+
+    public String getModel() {
+        return model;
+    }
+
+    public void setModel(String model) {
+        this.model = model;
+    }
+
+    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 Course getCourse() {
+        return course;
+    }
+
+    public void setCourse(Course course) {
+        this.course = course;
+    }
+
+    public Exercise getExercise() {
+        return exercise;
+    }
+
+    public void setExercise(Exercise exercise) {
+        this.exercise = exercise;
+    }
+
+    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 String getTraceId() {
+        return traceId;
+    }
+
+    public void setTraceId(String traceId) {
+        this.traceId = traceId;
+    }
+
+    public IrisMessage getIrisMessage() {
+        return irisMessage;
+    }
+
+    public void setIrisMessage(IrisMessage message) {
+        this.irisMessage = message;
+    }
+
+    @Override
+    public String toString() {
+        return "LLMTokenUsage{" + "serviceType=" + serviceType + ", model=" + model + ", numInputTokens=" + numInputTokens + ", costPerMillionInputTokens="
+                + costPerMillionInputTokens + ", numOutputTokens=" + numOutputTokens + ", costPerMillionOutputTokens=" + costPerMillionOutputTokens + ", course=" + course
+                + ", exercise=" + exercise + ", userId=" + userId + ", timestamp=" + time + ", traceId=" + traceId + ", irisMessage=" + irisMessage + '}';
+    }
+}

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

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

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

@@ -0,0 +1,110 @@
+package de.tum.cit.aet.artemis.core.service;
+
+import static de.tum.cit.aet.artemis.core.config.Constants.PROFILE_IRIS;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.context.annotation.Profile;
+import org.springframework.stereotype.Service;
+
+import de.tum.cit.aet.artemis.core.domain.Course;
+import de.tum.cit.aet.artemis.core.domain.LLMTokenUsage;
+import de.tum.cit.aet.artemis.core.domain.User;
+import de.tum.cit.aet.artemis.core.repository.LLMTokenUsageRepository;
+import de.tum.cit.aet.artemis.exercise.domain.Exercise;
+import de.tum.cit.aet.artemis.iris.domain.message.IrisMessage;
+import de.tum.cit.aet.artemis.iris.service.pyris.dto.data.PyrisLLMCostDTO;
+import de.tum.cit.aet.artemis.iris.service.pyris.job.PyrisJob;
+
+/**
+ * Service for managing the LLMTokenUsage by all LLMs in Artemis
+ */
+@Service
+@Profile(PROFILE_IRIS)
+public class LLMTokenUsageService {
+
+    private final LLMTokenUsageRepository llmTokenUsageRepository;
+
+    public LLMTokenUsageService(LLMTokenUsageRepository llmTokenUsageRepository) {
+        this.llmTokenUsageRepository = llmTokenUsageRepository;
+    }
+
+    /**
+     * method saves the token usage to the database with a link to the IrisMessage
+     * messages of the same job are grouped together by saving the job id as a trace id
+     *
+     * @param job      used to create a unique traceId to group multiple LLM calls
+     * @param message  IrisMessage to map the usage to an IrisMessage
+     * @param exercise to map the token cost to an exercise
+     * @param user     to map the token cost to a user
+     * @param course   to map the token to a course
+     * @param tokens   token cost list of type PyrisLLMCostDTO
+     * @return list of the saved data
+     */
+    public List<LLMTokenUsage> saveIrisTokenUsage(PyrisJob job, IrisMessage message, Exercise exercise, User user, Course course, List<PyrisLLMCostDTO> tokens) {
+        List<LLMTokenUsage> tokenUsages = new ArrayList<>();
+
+        for (PyrisLLMCostDTO cost : tokens) {
+            LLMTokenUsage llmTokenUsage = new LLMTokenUsage();
+            if (message != null) {
+                llmTokenUsage.setIrisMessage(message);
+                llmTokenUsage.setTime(message.getSentAt());
+            }
+            llmTokenUsage.setServiceType(cost.pipeline());
+            llmTokenUsage.setExercise(exercise);
+            if (user != null) {
+                llmTokenUsage.setUserId(user.getId());
+            }
+            llmTokenUsage.setCourse(course);
+            llmTokenUsage.setNumInputTokens(cost.numInputTokens());
+            llmTokenUsage.setCostPerMillionInputTokens(cost.costPerInputToken());
+            llmTokenUsage.setNumOutputTokens(cost.numOutputTokens());
+            llmTokenUsage.setCostPerMillionOutputTokens(cost.costPerOutputToken());
+            llmTokenUsage.setModel(cost.modelInfo());
+            llmTokenUsage.setTraceId(job.jobId());
+            tokenUsages.add(llmTokenUsage);
+        }
+        llmTokenUsageRepository.saveAll(tokenUsages);
+        return tokenUsages;
+    }
+
+    /**
+     * Overloaded method to save token usage without message and exercise.
+     *
+     * @param job    used to create a unique traceId to group multiple LLM calls
+     * @param user   to map the token cost to a user
+     * @param course to map the token to a course
+     * @param tokens token cost list of type PyrisLLMCostDTO
+     * @return list of the saved data
+     */
+    public List<LLMTokenUsage> saveIrisTokenUsage(PyrisJob job, User user, Course course, List<PyrisLLMCostDTO> tokens) {
+        return saveIrisTokenUsage(job, null, null, user, course, tokens);
+    }
+
+    /**
+     * Overloaded method to save token usage without exercise.
+     *
+     * @param job     used to create a unique traceId to group multiple LLM calls
+     * @param message IrisMessage to map the usage to an IrisMessage
+     * @param user    to map the token cost to a user
+     * @param course  to map the token to a course
+     * @param tokens  token cost list of type PyrisLLMCostDTO
+     * @return list of the saved data
+     */
+    public List<LLMTokenUsage> saveIrisTokenUsage(PyrisJob job, IrisMessage message, User user, Course course, List<PyrisLLMCostDTO> tokens) {
+        return saveIrisTokenUsage(job, message, null, user, course, tokens);
+    }
+
+    /**
+     * Overloaded method to save token usage without message, exercise and user.
+     *
+     * @param job    used to create a unique traceId to group multiple LLM calls
+     * @param course to map the token to a course
+     * @param tokens token cost list of type PyrisLLMCostDTO
+     * @return list of the saved data
+     */
+    public List<LLMTokenUsage> saveIrisTokenUsage(PyrisJob job, Course course, List<PyrisLLMCostDTO> tokens) {
+        return saveIrisTokenUsage(job, null, null, null, course, tokens);
+    }
+}

src/main/java/de/tum/cit/aet/artemis/iris/dto/IrisChatWebsocketDTO.java

@@ -9,6 +9,7 @@
 
 import de.tum.cit.aet.artemis.iris.domain.message.IrisMessage;
 import de.tum.cit.aet.artemis.iris.service.IrisRateLimitService;
+import de.tum.cit.aet.artemis.iris.service.pyris.dto.data.PyrisLLMCostDTO;
 import de.tum.cit.aet.artemis.iris.service.pyris.dto.status.PyrisStageDTO;
 
 /**
@@ -21,7 +22,7 @@
  */
 @JsonInclude(JsonInclude.Include.NON_EMPTY)
 public record IrisChatWebsocketDTO(IrisWebsocketMessageType type, IrisMessage message, IrisRateLimitService.IrisRateLimitInformation rateLimitInfo, List<PyrisStageDTO> stages,
-        List<String> suggestions) {
+        List<String> suggestions, List<PyrisLLMCostDTO> tokens) {
 
     /**
      * Creates a new IrisWebsocketDTO instance with the given parameters
@@ -31,8 +32,9 @@ public record IrisChatWebsocketDTO(IrisWebsocketMessageType type, IrisMessage me
      * @param rateLimitInfo the rate limit information
      * @param stages        the stages of the Pyris pipeline
      */
-    public IrisChatWebsocketDTO(@Nullable IrisMessage message, IrisRateLimitService.IrisRateLimitInformation rateLimitInfo, List<PyrisStageDTO> stages, List<String> suggestions) {
-        this(determineType(message), message, rateLimitInfo, stages, suggestions);
+    public IrisChatWebsocketDTO(@Nullable IrisMessage message, IrisRateLimitService.IrisRateLimitInformation rateLimitInfo, List<PyrisStageDTO> stages, List<String> suggestions,
+            List<PyrisLLMCostDTO> tokens) {
+        this(determineType(message), message, rateLimitInfo, stages, suggestions, tokens);
     }
 
     /**