-
Notifications
You must be signed in to change notification settings - Fork 291
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Iris
: Track token usage of iris requests
#9455
base: develop
Are you sure you want to change the base?
Changes from 18 commits
a5dfdbb
d0bdae3
2a08cb2
f85cf46
65fb259
188ff22
be85a3b
e974d59
6337162
84a60dc
5b0ab48
62dad8b
897d643
1d10860
8b27861
86294c1
56b20e7
8a29c82
abbd28f
8d34428
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
} |
Original file line number | Diff line number | Diff line change | ||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,162 @@ | ||||||||||||||||||
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.Table; | ||||||||||||||||||
|
||||||||||||||||||
import org.hibernate.annotations.Cache; | ||||||||||||||||||
import org.hibernate.annotations.CacheConcurrencyStrategy; | ||||||||||||||||||
|
||||||||||||||||||
import com.fasterxml.jackson.annotation.JsonInclude; | ||||||||||||||||||
|
||||||||||||||||||
@Entity | ||||||||||||||||||
@Table(name = "llm_token_usage") | ||||||||||||||||||
@Cache(usage = CacheConcurrencyStrategy.NONSTRICT_READ_WRITE) | ||||||||||||||||||
@JsonInclude(JsonInclude.Include.NON_EMPTY) | ||||||||||||||||||
public class LLMTokenUsage extends DomainObject { | ||||||||||||||||||
alexjoham marked this conversation as resolved.
Show resolved
Hide resolved
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧹 Nitpick (assertive) Consider adding JavaDoc comments for the class Adding a class-level JavaDoc comment can improve code readability and help other developers understand the purpose and usage of this entity. |
||||||||||||||||||
|
||||||||||||||||||
@Column(name = "service") | ||||||||||||||||||
private String serviceType; | ||||||||||||||||||
alexjoham marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||
|
||||||||||||||||||
@Column(name = "model") | ||||||||||||||||||
private String model; | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion Consider Using an Enum for Similar to Apply these changes:
|
||||||||||||||||||
|
||||||||||||||||||
@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; | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Use Using Apply these changes:
|
||||||||||||||||||
|
||||||||||||||||||
@Nullable | ||||||||||||||||||
@Column(name = "course_id") | ||||||||||||||||||
private Long courseId; | ||||||||||||||||||
|
||||||||||||||||||
@Nullable | ||||||||||||||||||
@Column(name = "exercise_id") | ||||||||||||||||||
private Long exerciseId; | ||||||||||||||||||
MichaelOwenDyer marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||
|
||||||||||||||||||
@Column(name = "user_id") | ||||||||||||||||||
private long userId; | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion Consider Mapping Mapping the Apply these changes:
|
||||||||||||||||||
|
||||||||||||||||||
@Column(name = "time") | ||||||||||||||||||
private ZonedDateTime time = ZonedDateTime.now(); | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Avoid using reserved SQL keywords for column names The column name Apply this change: - @Column(name = "time")
+ @Column(name = "\"time\"")
private ZonedDateTime time = ZonedDateTime.now(); Or rename the column: - @Column(name = "time")
+ @Column(name = "timestamp")
private ZonedDateTime time = ZonedDateTime.now(); Ensure that any database scripts or queries are updated to reflect this change. 📝 Committable suggestion
Suggested change
|
||||||||||||||||||
|
||||||||||||||||||
@Column(name = "trace_id") | ||||||||||||||||||
private String traceId; | ||||||||||||||||||
|
||||||||||||||||||
@Nullable | ||||||||||||||||||
@Column(name = "iris_message_id") | ||||||||||||||||||
private Long irisMessageId; | ||||||||||||||||||
|
||||||||||||||||||
public String getServiceType() { | ||||||||||||||||||
return serviceType; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public void setServiceType(String 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; | ||||||||||||||||||
} | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧹 Nitpick (assertive) Synchronize Parameter Names in Setter Methods The parameter name in Apply this diff: - public void setCostPerMillionInputTokens(float costPerMillionInputToken) {
- this.costPerMillionInputTokens = costPerMillionInputToken;
+ public void setCostPerMillionInputTokens(float costPerMillionInputTokens) {
+ this.costPerMillionInputTokens = costPerMillionInputTokens;
} Similarly, update the - public void setCostPerMillionOutputTokens(float costPerMillionOutputToken) {
- this.costPerMillionOutputTokens = costPerMillionOutputToken;
+ public void setCostPerMillionOutputTokens(float costPerMillionOutputTokens) {
+ this.costPerMillionOutputTokens = costPerMillionOutputTokens;
} 📝 Committable suggestion
Suggested change
|
||||||||||||||||||
|
||||||||||||||||||
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 Long getCourseId() { | ||||||||||||||||||
return courseId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public void setCourseId(Long courseId) { | ||||||||||||||||||
this.courseId = courseId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public Long getExercisIde() { | ||||||||||||||||||
return exerciseId; | ||||||||||||||||||
} | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fix Typo in Method Name The getter method for Apply this diff to correct the method name: - public Long getExercisIde() {
+ public Long getExerciseId() {
return exerciseId;
} 📝 Committable suggestion
Suggested change
|
||||||||||||||||||
|
||||||||||||||||||
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 String getTraceId() { | ||||||||||||||||||
return traceId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public void setTraceId(String traceId) { | ||||||||||||||||||
this.traceId = traceId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public Long getIrisMessageId() { | ||||||||||||||||||
return irisMessageId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
public void setIrisMessageId(Long messageId) { | ||||||||||||||||||
this.irisMessageId = messageId; | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
@Override | ||||||||||||||||||
public String toString() { | ||||||||||||||||||
return "LLMTokenUsage{" + "serviceType=" + serviceType + ", model=" + model + ", numInputTokens=" + numInputTokens + ", costPerMillionInputTokens=" | ||||||||||||||||||
+ costPerMillionInputTokens + ", numOutputTokens=" + numOutputTokens + ", costPerMillionOutputTokens=" + costPerMillionOutputTokens + ", course=" + courseId | ||||||||||||||||||
+ ", exercise=" + exerciseId + ", userId=" + userId + ", timestamp=" + time + ", traceId=" + traceId + ", irisMessage=" + irisMessageId + '}'; | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧹 Nitpick (assertive) Ensure Consistency in In the Apply this diff: + ", exercise=" + exerciseId + ", userId=" + userId + ",
- timestamp=" + time + ", traceId=" + traceId + ", irisMessage=" + irisMessageId + '}';
+ time=" + time + ", traceId=" + traceId + ", irisMessage=" + irisMessageId + '}'; 📝 Committable suggestion
Suggested change
|
||||||||||||||||||
} | ||||||||||||||||||
} | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧹 Nitpick (assertive) Add Class-Level JavaDoc Comment Adding a JavaDoc comment to the class Apply this change: /**
* Entity representing the usage of tokens by Large Language Models (LLMs).
* It tracks input/output token counts, associated costs, and related metadata
* such as the user, course, exercise, and timestamps.
*/
public class LLMTokenUsage extends DomainObject { |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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> { | ||
} |
Original file line number | Diff line number | Diff line change | ||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,151 @@ | ||||||||||||
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 java.util.Optional; | ||||||||||||
import java.util.function.Function; | ||||||||||||
|
||||||||||||
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 builderFunction of type Function<IrisTokenUsageBuilder, IrisTokenUsageBuilder> using IrisTokenUsageBuilder | ||||||||||||
* @return saved LLMTokenUsage as a List | ||||||||||||
*/ | ||||||||||||
public List<LLMTokenUsage> saveIrisTokenUsage(Function<IrisTokenUsageBuilder, IrisTokenUsageBuilder> builderFunction) { | ||||||||||||
IrisTokenUsageBuilder builder = builderFunction.apply(new IrisTokenUsageBuilder()); | ||||||||||||
List<LLMTokenUsage> tokenUsages = new ArrayList<>(); | ||||||||||||
List<PyrisLLMCostDTO> tokens = builder.getTokens(); | ||||||||||||
for (PyrisLLMCostDTO cost : tokens) { | ||||||||||||
LLMTokenUsage llmTokenUsage = new LLMTokenUsage(); | ||||||||||||
|
||||||||||||
builder.getMessage().ifPresent(message -> { | ||||||||||||
llmTokenUsage.setIrisMessageId(message.getId()); | ||||||||||||
llmTokenUsage.setTime(message.getSentAt()); | ||||||||||||
}); | ||||||||||||
|
||||||||||||
builder.getUser().ifPresent(user -> { | ||||||||||||
llmTokenUsage.setUserId(user.getId()); | ||||||||||||
}); | ||||||||||||
|
||||||||||||
builder.getExercise().ifPresent(exercise -> { | ||||||||||||
llmTokenUsage.setExerciseId(exercise.getId()); | ||||||||||||
}); | ||||||||||||
|
||||||||||||
builder.getCourse().ifPresent(course -> { | ||||||||||||
llmTokenUsage.setCourseId(course.getId()); | ||||||||||||
}); | ||||||||||||
|
||||||||||||
llmTokenUsage.setTraceId(builder.getJob().jobId()); | ||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Add null check for 'job' before using it In the Apply this diff to add the null check: + if (builder.getJob() == null) {
+ throw new IllegalArgumentException("Job must not be null");
+ }
llmTokenUsage.setTraceId(builder.getJob().jobId()); 📝 Committable suggestion
Suggested change
|
||||||||||||
llmTokenUsage.setServiceType(cost.pipeline()); | ||||||||||||
llmTokenUsage.setNumInputTokens(cost.numInputTokens()); | ||||||||||||
llmTokenUsage.setCostPerMillionInputTokens(cost.costPerInputToken()); | ||||||||||||
llmTokenUsage.setNumOutputTokens(cost.numOutputTokens()); | ||||||||||||
llmTokenUsage.setCostPerMillionOutputTokens(cost.costPerOutputToken()); | ||||||||||||
llmTokenUsage.setModel(cost.modelInfo()); | ||||||||||||
tokenUsages.add(llmTokenUsage); | ||||||||||||
} | ||||||||||||
llmTokenUsageRepository.saveAll(tokenUsages); | ||||||||||||
return tokenUsages; | ||||||||||||
} | ||||||||||||
|
||||||||||||
/** | ||||||||||||
* Class IrisTokenUsageBuilder to be used for saveIrisTokenUsage() | ||||||||||||
*/ | ||||||||||||
public static class IrisTokenUsageBuilder { | ||||||||||||
|
||||||||||||
private PyrisJob job; | ||||||||||||
|
||||||||||||
private List<PyrisLLMCostDTO> tokens; | ||||||||||||
|
||||||||||||
private Optional<Course> course = Optional.empty(); | ||||||||||||
|
||||||||||||
private Optional<IrisMessage> message = Optional.empty(); | ||||||||||||
|
||||||||||||
private Optional<Exercise> exercise = Optional.empty(); | ||||||||||||
|
||||||||||||
private Optional<User> user = Optional.empty(); | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withJob(PyrisJob job) { | ||||||||||||
this.job = job; | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withCourse(Course course) { | ||||||||||||
this.course = Optional.ofNullable(course); | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withTokens(List<PyrisLLMCostDTO> tokens) { | ||||||||||||
this.tokens = tokens; | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withMessage(IrisMessage message) { | ||||||||||||
this.message = Optional.ofNullable(message); | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withExercise(Exercise exercise) { | ||||||||||||
this.exercise = Optional.ofNullable(exercise); | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public IrisTokenUsageBuilder withUser(User user) { | ||||||||||||
this.user = Optional.ofNullable(user); | ||||||||||||
return this; | ||||||||||||
} | ||||||||||||
|
||||||||||||
// Getters | ||||||||||||
public PyrisJob getJob() { | ||||||||||||
return job; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public List<PyrisLLMCostDTO> getTokens() { | ||||||||||||
return tokens; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public Optional<Course> getCourse() { | ||||||||||||
return course; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public Optional<IrisMessage> getMessage() { | ||||||||||||
return message; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public Optional<Exercise> getExercise() { | ||||||||||||
return exercise; | ||||||||||||
} | ||||||||||||
|
||||||||||||
public Optional<User> getUser() { | ||||||||||||
return user; | ||||||||||||
} | ||||||||||||
} | ||||||||||||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🧹 Nitpick (assertive)
Enhance JavaDoc with details about enum constants.
The JavaDoc comment provides a good overview of the enum's purpose. However, consider adding more details about the enum constants to improve documentation.
Here's a suggested enhancement: