Iris: Track token usage of iris requests

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

@@ -0,0 +1,6 @@
+package de.tum.cit.aet.artemis.core.domain;
+
+public enum LLMServiceType {
+    ATHENA, IRIS_CODE_FEEDBACK, IRIS_CHAT_COURSE_MESSAGE, IRIS_CHAT_EXERCISE_MESSAGE, IRIS_INTERACTION_SUGGESTION, IRIS_CHAT_LECTURE_MESSAGE, IRIS_COMPETENCY_GENERATION,
+    IRIS_CITATION_PIPELINE, NOT_SET
+}

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

@@ -0,0 +1,154 @@
+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.Inheritance;
+import jakarta.persistence.InheritanceType;
+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 com.fasterxml.jackson.annotation.JsonTypeInfo;
+
+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")
+@Inheritance(strategy = InheritanceType.SINGLE_TABLE)
+@Cache(usage = CacheConcurrencyStrategy.NONSTRICT_READ_WRITE)
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
+@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 = "cost_per_token")
+    private double cost_per_token;
+
+    @Column(name = "num_input_tokens")
+    private int num_input_tokens;
+
+    @Column(name = "num_output_tokens")
+    private int num_output_tokens;
+
+    @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;
+
+    @Nullable
+    @Column(name = "timestamp")
+    private ZonedDateTime timestamp = ZonedDateTime.now();
+
+    @Nullable
+    @ManyToOne
+    @JsonIgnore
+    @JoinColumn(name = "iris_message_id")
+    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 double getCost_per_token() {
+        return cost_per_token;
+    }
+
+    public void setCost_per_token(double cost_per_token) {
+        this.cost_per_token = cost_per_token;
+    }
+
+    public int getNum_input_tokens() {
+        return num_input_tokens;
+    }
+
+    public void setNum_input_tokens(int num_input_tokens) {
+        this.num_input_tokens = num_input_tokens;
+    }
+
+    public int getNum_output_tokens() {
+        return num_output_tokens;
+    }
+
+    public void setNum_output_tokens(int num_output_tokens) {
+        this.num_output_tokens = num_output_tokens;
+    }
+
+    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 getTimestamp() {
+        return timestamp;
+    }
+
+    public void setTimestamp(ZonedDateTime timestamp) {
+        this.timestamp = timestamp;
+    }
+
+    public IrisMessage getIrisMessage() {
+        return irisMessage;
+    }
+
+    public void setIrisMessage(IrisMessage message) {
+        this.irisMessage = message;
+    }
+}

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

@@ -0,0 +1,10 @@
+package de.tum.cit.aet.artemis.core.repository;
+
+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
+public interface LLMTokenUsageRepository extends ArtemisJpaRepository<LLMTokenUsage, Long> {
+}

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

@@ -0,0 +1,66 @@
+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;
+
+/**
+ * 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;
+    }
+
+    /**
+     * saves the tokens used for a specific IrisMessage or Athena call
+     * in case of an Athena call IrisMessage can be null and the
+     * LLMServiceType in tokens has to by Athena
+     *
+     * @param message  IrisMessage related to the TokenUsage
+     * @param exercise Exercise in which the request was made
+     * @param user     User that made the request
+     * @param course   Course in which the request was made
+     * @param tokens   List with Tokens of the PyrisLLMCostDTO Mdel
+     * @return List of the created LLMTokenUsage entries
+     */
+
+    public List<LLMTokenUsage> saveTokenUsage(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.setTimestamp(message.getSentAt());
+            }
+            llmTokenUsage.setServiceType(cost.pipeline());
+            llmTokenUsage.setExercise(exercise);
+            if (user != null) {
+                llmTokenUsage.setUserId(user.getId());
+            }
+            llmTokenUsage.setCourse(course);
+            llmTokenUsage.setNum_input_tokens(cost.num_input_tokens());
+            llmTokenUsage.setNum_output_tokens(cost.num_output_tokens());
+            llmTokenUsage.setModel(cost.model_info());
+            tokenUsages.add(llmTokenUsageRepository.save(llmTokenUsage));
+        }
+        return tokenUsages;
+    }
+}

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);
     }
 
     /**

src/main/java/de/tum/cit/aet/artemis/iris/service/IrisCompetencyGenerationService.java

@@ -8,6 +8,7 @@
 import de.tum.cit.aet.artemis.atlas.domain.competency.CompetencyTaxonomy;
 import de.tum.cit.aet.artemis.core.domain.Course;
 import de.tum.cit.aet.artemis.core.domain.User;
+import de.tum.cit.aet.artemis.core.service.LLMTokenUsageService;
 import de.tum.cit.aet.artemis.iris.service.pyris.PyrisJobService;
 import de.tum.cit.aet.artemis.iris.service.pyris.PyrisPipelineService;
 import de.tum.cit.aet.artemis.iris.service.pyris.dto.competency.PyrisCompetencyExtractionPipelineExecutionDTO;
@@ -25,14 +26,18 @@ public class IrisCompetencyGenerationService {
 
     private final PyrisPipelineService pyrisPipelineService;
 
+    private final LLMTokenUsageService llmTokenUsageService;
+
     private final IrisWebsocketService websocketService;
 
     private final PyrisJobService pyrisJobService;
 
-    public IrisCompetencyGenerationService(PyrisPipelineService pyrisPipelineService, IrisWebsocketService websocketService, PyrisJobService pyrisJobService) {
+    public IrisCompetencyGenerationService(PyrisPipelineService pyrisPipelineService, LLMTokenUsageService llmTokenUsageService, IrisWebsocketService websocketService,
+            PyrisJobService pyrisJobService) {
         this.pyrisPipelineService = pyrisPipelineService;
         this.websocketService = websocketService;
         this.pyrisJobService = pyrisJobService;
+        this.llmTokenUsageService = llmTokenUsageService;
     }
 
     /**
@@ -50,7 +55,7 @@ public void executeCompetencyExtractionPipeline(User user, Course course, String
                 "default",
                 pyrisJobService.createTokenForJob(token -> new CompetencyExtractionJob(token, course.getId(), user.getLogin())),
                 executionDto -> new PyrisCompetencyExtractionPipelineExecutionDTO(executionDto, courseDescription, currentCompetencies, CompetencyTaxonomy.values(), 5),
-                stages -> websocketService.send(user.getLogin(), websocketTopic(course.getId()), new PyrisCompetencyStatusUpdateDTO(stages, null))
+                stages -> websocketService.send(user.getLogin(), websocketTopic(course.getId()), new PyrisCompetencyStatusUpdateDTO(stages, null, null))
         );
         // @formatter:on
     }
@@ -63,6 +68,7 @@ public void executeCompetencyExtractionPipeline(User user, Course course, String
      * @param statusUpdate the status update containing the new competency recommendations
      */
     public void handleStatusUpdate(String userLogin, long courseId, PyrisCompetencyStatusUpdateDTO statusUpdate) {
+        var tokenUsages = llmTokenUsageService.saveTokenUsage(null, null, null, null, statusUpdate.tokens());
         websocketService.send(userLogin, websocketTopic(courseId), statusUpdate);
     }
 

src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/dto/chat/PyrisChatStatusUpdateDTO.java

@@ -4,8 +4,9 @@
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 
+import de.tum.cit.aet.artemis.iris.service.pyris.dto.data.PyrisLLMCostDTO;
 import de.tum.cit.aet.artemis.iris.service.pyris.dto.status.PyrisStageDTO;
 
 @JsonInclude(JsonInclude.Include.NON_EMPTY)
-public record PyrisChatStatusUpdateDTO(String result, List<PyrisStageDTO> stages, List<String> suggestions) {
+public record PyrisChatStatusUpdateDTO(String result, List<PyrisStageDTO> stages, List<String> suggestions, List<PyrisLLMCostDTO> tokens) {
 }

src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/dto/competency/PyrisCompetencyStatusUpdateDTO.java

@@ -4,6 +4,7 @@
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 
+import de.tum.cit.aet.artemis.iris.service.pyris.dto.data.PyrisLLMCostDTO;
 import de.tum.cit.aet.artemis.iris.service.pyris.dto.status.PyrisStageDTO;
 
 /**
@@ -15,5 +16,5 @@
  * @param result List of competencies recommendations that have been generated so far
  */
 @JsonInclude(JsonInclude.Include.NON_EMPTY)
-public record PyrisCompetencyStatusUpdateDTO(List<PyrisStageDTO> stages, List<PyrisCompetencyRecommendationDTO> result) {
+public record PyrisCompetencyStatusUpdateDTO(List<PyrisStageDTO> stages, List<PyrisCompetencyRecommendationDTO> result, List<PyrisLLMCostDTO> tokens) {
 }

src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/dto/data/PyrisLLMCostDTO.java

@@ -0,0 +1,6 @@
+package de.tum.cit.aet.artemis.iris.service.pyris.dto.data;
+
+import de.tum.cit.aet.artemis.core.domain.LLMServiceType;
+
+public record PyrisLLMCostDTO(String model_info, int num_input_tokens, int num_output_tokens, LLMServiceType pipeline) {
+}

src/main/java/de/tum/cit/aet/artemis/iris/service/session/IrisCourseChatSessionService.java

@@ -19,6 +19,7 @@
 import de.tum.cit.aet.artemis.core.exception.AccessForbiddenException;
 import de.tum.cit.aet.artemis.core.security.Role;
 import de.tum.cit.aet.artemis.core.service.AuthorizationCheckService;
+import de.tum.cit.aet.artemis.core.service.LLMTokenUsageService;
 import de.tum.cit.aet.artemis.iris.domain.message.IrisMessage;
 import de.tum.cit.aet.artemis.iris.domain.message.IrisMessageSender;
 import de.tum.cit.aet.artemis.iris.domain.message.IrisTextMessageContent;
@@ -43,6 +44,8 @@ public class IrisCourseChatSessionService extends AbstractIrisChatSessionService
 
     private final IrisMessageService irisMessageService;
 
+    private final LLMTokenUsageService llmTokenUsageService;
+
     private final IrisSettingsService irisSettingsService;
 
     private final IrisChatWebsocketService irisChatWebsocketService;
@@ -57,11 +60,13 @@ public class IrisCourseChatSessionService extends AbstractIrisChatSessionService
 
     private final PyrisPipelineService pyrisPipelineService;
 
-    public IrisCourseChatSessionService(IrisMessageService irisMessageService, IrisSettingsService irisSettingsService, IrisChatWebsocketService irisChatWebsocketService,
-            AuthorizationCheckService authCheckService, IrisSessionRepository irisSessionRepository, IrisRateLimitService rateLimitService,
-            IrisCourseChatSessionRepository irisCourseChatSessionRepository, PyrisPipelineService pyrisPipelineService, ObjectMapper objectMapper) {
+    public IrisCourseChatSessionService(IrisMessageService irisMessageService, LLMTokenUsageService llmTokenUsageService, IrisSettingsService irisSettingsService,
+            IrisChatWebsocketService irisChatWebsocketService, AuthorizationCheckService authCheckService, IrisSessionRepository irisSessionRepository,
+            IrisRateLimitService rateLimitService, IrisCourseChatSessionRepository irisCourseChatSessionRepository, PyrisPipelineService pyrisPipelineService,
+            ObjectMapper objectMapper, LLMTokenUsageService lLMTokenUsageService) {
         super(irisSessionRepository, objectMapper);
         this.irisMessageService = irisMessageService;
+        this.llmTokenUsageService = llmTokenUsageService;
         this.irisSettingsService = irisSettingsService;
         this.irisChatWebsocketService = irisChatWebsocketService;
         this.authCheckService = authCheckService;
@@ -137,10 +142,12 @@ public void handleStatusUpdate(CourseChatJob job, PyrisChatStatusUpdateDTO statu
             var message = new IrisMessage();
             message.addContent(new IrisTextMessageContent(statusUpdate.result()));
             var savedMessage = irisMessageService.saveMessage(message, session, IrisMessageSender.LLM);
+            var tokenUsages = llmTokenUsageService.saveTokenUsage(savedMessage, null, session.getUser(), session.getCourse(), statusUpdate.tokens());
             irisChatWebsocketService.sendMessage(session, savedMessage, statusUpdate.stages());
         }
         else {
-            irisChatWebsocketService.sendStatusUpdate(session, statusUpdate.stages(), statusUpdate.suggestions());
+            var tokenUsages = llmTokenUsageService.saveTokenUsage(null, null, session.getUser(), session.getCourse(), statusUpdate.tokens());
+            irisChatWebsocketService.sendStatusUpdate(session, statusUpdate.stages(), statusUpdate.suggestions(), statusUpdate.tokens());
         }
         updateLatestSuggestions(session, statusUpdate.suggestions());
     }