Athena: Add LLM token usage tracking
#9554
Conversation
github-actions
bot
added
tests
server
FelixTJDietrich
changed the base branch from
develop
to
feature/track-usage-of-iris-requests
WalkthroughThis pull request introduces several new classes and modifications to existing services and DTOs, primarily to enhance the handling of token usage for large language models (LLMs) within the Artemis system. Key additions include the Changes
Possibly related PRs
Suggested labels
Suggested reviewers
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's limits when posting large numbers of comments.
Actionable comments posted: 30
🧹 Outside diff range and nitpick comments (22)
src/main/java/de/tum/cit/aet/artemis/core/domain/LLMServiceType.java (1)
1-8: LGTM! Consider enhancing the documentation.The
LLMServiceTypeenum is well-structured and follows the coding guidelines. It adheres to the single responsibility principle and uses appropriate naming conventions. The class-level comment provides a clear description of the enum's purpose.Consider adding brief descriptions for each enum constant to provide more context:
public enum LLMServiceType { /** Represents the IRIS LLM service. */ IRIS, /** Represents the ATHENA LLM service. */ ATHENA }This addition would enhance the documentation and make it easier for developers to understand the specific purpose of each service type.
src/main/java/de/tum/cit/aet/artemis/core/domain/LLMRequest.java (1)
1-4: Add Javadoc comments to improve documentation.To enhance code readability and maintainability, consider adding Javadoc comments for the
LLMRequestrecord and its fields. This will provide clear documentation for other developers who might use this record.Here's a suggested Javadoc comment structure:
/** * Represents a request for a Large Language Model (LLM) operation. * * @param model The name or identifier of the LLM model. * @param numInputTokens The number of input tokens for the request. * @param costPerMillionInputToken The cost per million input tokens. * @param numOutputTokens The number of output tokens generated. * @param costPerMillionOutputToken The cost per million output tokens. * @param pipelineId The identifier of the pipeline used for this request. */ public record LLMRequest(String model, int numInputTokens, BigDecimal costPerMillionInputToken, int numOutputTokens, BigDecimal costPerMillionOutputToken, String pipelineId) { }src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/dto/data/PyrisLLMCostDTO.java (1)
1-4: LGTM! Consider adding Javadoc for better documentation.The
PyrisLLMCostDTOrecord is well-structured and follows the coding guidelines. It uses appropriate types for its fields and adheres to the single responsibility principle by focusing solely on LLM cost data.Consider adding Javadoc to the record and its fields for improved documentation. Here's a suggestion:
/** * Data Transfer Object for PyrisLLM cost information. */ public record PyrisLLMCostDTO( /** Information about the model used. */ String modelInfo, /** Number of input tokens. */ int numInputTokens, /** Cost per input token. */ float costPerInputToken, /** Number of output tokens. */ int numOutputTokens, /** Cost per output token. */ float costPerOutputToken, /** Information about the pipeline used. */ String pipeline ) {}src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/job/SessionBasedPyrisJob.java (1)
3-5: LGTM: Clear and concise Javadoc comment.The Javadoc comment effectively describes the purpose of the interface. Consider adding a brief explanation of what a "session" means in this context for improved clarity.
src/main/java/de/tum/cit/aet/artemis/core/repository/LLMTokenUsageTraceRepository.java (1)
11-14: LGTM: Interface declaration and annotations are correct.The interface is well-defined and follows the coding guidelines. It correctly extends ArtemisJpaRepository and is properly annotated. The empty interface is acceptable as it inherits standard CRUD operations.
Consider adding a brief Javadoc comment to describe the purpose of this repository interface. For example:
/** * Repository interface for managing LLMTokenUsageTrace entities. * Provides CRUD operations and custom queries if needed. */ @Profile(PROFILE_CORE) @Repository public interface LLMTokenUsageTraceRepository extends ArtemisJpaRepository<LLMTokenUsageTrace, Long> { }src/main/java/de/tum/cit/aet/artemis/core/repository/LLMTokenUsageRequestRepository.java (1)
11-13: LGTM: Interface declaration and annotations are correct.The interface is well-defined and follows the correct naming convention. The annotations
@Profile(PROFILE_CORE)and@Repositoryare appropriately used. ExtendingArtemisJpaRepositorywith the correct entity type and ID type is good.Consider adding a brief Javadoc comment to describe the purpose of this repository interface. For example:
/** * Repository interface for managing LLMTokenUsageRequest entities. * This interface provides CRUD operations for LLMTokenUsageRequest objects. */ @Profile(PROFILE_CORE) @Repository public interface LLMTokenUsageRequestRepository extends ArtemisJpaRepository<LLMTokenUsageRequest, Long> { }src/main/java/de/tum/cit/aet/artemis/athena/dto/ResponseMetaDTO.java (1)
13-14: LGTM: Well-structured DTO using Java record.The
ResponseMetaDTOrecord is well-defined and follows good practices for DTOs. It adheres to the single responsibility principle and uses appropriate types for its fields.Consider expanding the Javadoc comment to briefly describe the purpose of each field in the record. This would enhance the documentation and make it easier for other developers to understand the DTO's structure.
src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/job/CourseChatJob.java (1)
12-12: LGTM! Consider enhancing the documentation.The change from
PyrisJobtoSessionBasedPyrisJobis appropriate and aligns with the session-based nature of the job. The implementation is concise and follows good practices.Consider updating the class-level documentation to reflect the change to
SessionBasedPyrisJoband briefly explain the significance of thesessionIdfield. This would improve code clarity for future developers.src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/job/ExerciseChatJob.java (1)
13-13: LGTM! Consider adding a brief Javadoc for the new interface.The change from
PyrisJobtoSessionBasedPyrisJobaligns well with the PR objectives for LLM token usage tracking. The inclusion ofsessionIdin the record is consistent with this change.Consider adding a brief Javadoc comment for the
SessionBasedPyrisJobinterface to explain its purpose and how it differs from the generalPyrisJob. This would enhance code documentation and maintainability.src/main/java/de/tum/cit/aet/artemis/iris/service/pyris/dto/competency/PyrisCompetencyStatusUpdateDTO.java (2)
17-17: Javadoc updated correctly, minor grammatical fix needed.The new parameter description is clear and informative. However, there's a small grammatical error.
Please apply this minor correction:
- * @param tokens List of token usages send by Pyris for tracking the token usage and cost + * @param tokens List of token usages sent by Pyris for tracking the token usage and cost
20-20: Record updated correctly, consider adding a@Nonnullannotation.The addition of the
tokensparameter aligns with the PR objectives and follows the coding guidelines. The record is well-structured and uses appropriate annotations.For consistency and to prevent potential null pointer exceptions, consider adding a
@Nonnullannotation to thetokensparameter:- public record PyrisCompetencyStatusUpdateDTO(List<PyrisStageDTO> stages, List<PyrisCompetencyRecommendationDTO> result, List<LLMRequest> tokens) { + public record PyrisCompetencyStatusUpdateDTO(List<PyrisStageDTO> stages, List<PyrisCompetencyRecommendationDTO> result, @Nonnull List<LLMRequest> tokens) {This change would ensure that the
tokenslist is always non-null, which aligns with best practices for null safety.src/main/java/de/tum/cit/aet/artemis/iris/dto/IrisChatWebsocketDTO.java (2)
35-37: LGTM: Constructor updated correctly with a minor suggestion.The constructor changes are consistent with the record declaration update and maintain the existing functionality. They follow the coding guidelines by using constructor injection and keeping the code simple.
Consider breaking the long constructor parameter list into multiple lines for improved readability:
public IrisChatWebsocketDTO( @Nullable IrisMessage message, IrisRateLimitService.IrisRateLimitInformation rateLimitInfo, List<PyrisStageDTO> stages, List<String> suggestions, List<LLMRequest> tokens ) { this(determineType(message), message, rateLimitInfo, stages, suggestions, tokens); }
Line range hint
1-79: Summary: Changes align well with PR objectives and coding guidelines.The modifications to
IrisChatWebsocketDTOeffectively implement the tracking of LLM token usage as intended in the PR objectives. The changes maintain consistency with existing code structure and adhere to the project's coding guidelines. The addition of thetokensparameter enhances the DTO's capability to transfer LLM token usage data, which is crucial for the new feature being implemented.As the system evolves to include more LLM-related features, consider creating a separate DTO for LLM-specific data if the amount of data grows significantly. This would help maintain the single responsibility principle and keep the
IrisChatWebsocketDTOfocused on its primary purpose.src/test/java/de/tum/cit/aet/artemis/iris/IrisChatWebsocketTest.java (1)
56-56: Enhance test to verify the new parameter.While the test has been updated to include the new parameter in the
IrisChatWebsocketDTOconstructor, it doesn't explicitly verify this new field. Consider enhancing the test to ensure the new parameter is correctly handled.You could modify the test to include a non-empty list for the new parameter and verify that it's correctly passed to the
sendMessageToUsermethod. For example:List<LLMRequest> mockTokens = List.of(new LLMRequest("test", 10)); irisChatWebsocketService.sendMessage(irisSession, message, mockTokens); verify(websocketMessagingService, times(1)).sendMessageToUser(eq(TEST_PREFIX + "student1"), eq("/topic/iris/" + irisSession.getId()), eq(new IrisChatWebsocketDTO(message, new IrisRateLimitService.IrisRateLimitInformation(0, -1, 0), List.of(), List.of(), mockTokens)));This change would make the test more robust and ensure that the new functionality is correctly tested.
src/test/java/de/tum/cit/aet/artemis/iris/IrisCompetencyGenerationIntegrationTest.java (1)
70-73: LGTM: Test updated to reflect API changes.The test has been correctly updated to use the new
CompetencyExtractionJobobject, which aligns with the changes in theIrisCompetencyGenerationServiceAPI. The test remains focused and specific to the functionality being tested.Consider adding a comment explaining the purpose of the
nullargument in thePyrisCompetencyStatusUpdateDTOconstructor for better clarity.🧰 Tools
🪛 ast-grep
[warning] 72-72: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: (job, new PyrisCompetencyStatusUpdateDTO(stages, recommendations, null))
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 72-72: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: (job, new PyrisCompetencyStatusUpdateDTO(stages, recommendations, null))
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfigurationsrc/test/java/de/tum/cit/aet/artemis/iris/IrisTextExerciseChatMessageIntegrationTest.java (1)
401-401: LGTM. Consider adding a TODO comment for token usage.The addition of the
nullparameter toPyrisChatStatusUpdateDTOaligns with the PR objective of implementing LLM token usage tracking. However, it might be beneficial to add a TODO comment explaining why this parameter is currently null and when it will be implemented.Consider adding a TODO comment like this:
- request.postWithoutResponseBody("/api/public/pyris/pipelines/text-exercise-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, suggestions, null), + // TODO: Implement token usage tracking once the feature is fully developed + request.postWithoutResponseBody("/api/public/pyris/pipelines/text-exercise-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, suggestions, null), HttpStatus.OK, headers);🧰 Tools
🪛 ast-grep
[warning] 400-401: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: ("/api/public/pyris/pipelines/text-exercise-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, suggestions, null),
HttpStatus.OK, headers)
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 400-401: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: ("/api/public/pyris/pipelines/text-exercise-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, suggestions, null),
HttpStatus.OK, headers)
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfigurationsrc/main/java/de/tum/cit/aet/artemis/core/domain/LLMTokenUsageRequest.java (1)
3-11: Avoid unused imports and maintain import orderEnsure that all imported packages are necessary. Additionally, follow the standard import ordering: Java standard library imports, third-party library imports, and finally, project-specific imports.
Consider organizing the imports like this:
package de.tum.cit.aet.artemis.core.domain; +import com.fasterxml.jackson.annotation.JsonInclude; +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; // Remove any unused imports if presentsrc/main/java/de/tum/cit/aet/artemis/iris/service/websocket/IrisChatWebsocketService.java (1)
65-65: Correct grammatical error in Javadoc commentThe Javadoc for
@param tokensshould be "token usage and cost sent by Pyris" instead of "token usage and cost send by Pyris".src/main/java/de/tum/cit/aet/artemis/iris/service/IrisCompetencyGenerationService.java (1)
78-85: Ensure consistent use of explicit type declarationsIn line 84,
varis used for the variable declaration. For consistency and readability, consider using an explicit type declaration:-var user = userRepository.findById(job.userId()).orElseThrow(); +User user = userRepository.findById(job.userId()).orElseThrow();src/main/java/de/tum/cit/aet/artemis/iris/service/session/AbstractIrisChatSessionService.java (1)
87-113: Ensure proper management of entries intracesto prevent memory leaksEntries in the
tracesmap are added and removed based on the job ID. Ensure that all entries are removed when no longer needed to prevent the map from growing indefinitely, which could lead to memory leaks over time.src/main/java/de/tum/cit/aet/artemis/iris/service/session/IrisExerciseChatSessionService.java (1)
155-158: Add unit tests forsetLLMTokenUsageParametersmethodConsider adding unit tests for the new
setLLMTokenUsageParametersmethod to verify its functionality and handle potential edge cases.src/test/java/de/tum/cit/aet/artemis/iris/IrisChatTokenTrackingIntegrationTest.java (1)
225-229: Specify thrown exceptions for clarityThe method
sendStatusdeclares a generalthrows Exception, which can obscure the specific exceptions that might be thrown. For better clarity and exception handling, consider specifying the exact exceptions that can occur, such asIOExceptionorHttpClientErrorException.🧰 Tools
🪛 ast-grep
[warning] 225-225: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: (Map.of("Authorization", List.of("Bearer " + jobId)))
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 226-227: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: ("/api/public/pyris/pipelines/tutor-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, null, tokens),
HttpStatus.OK, headers)
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 225-225: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: (Map.of("Authorization", List.of("Bearer " + jobId)))
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 226-227: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: ("/api/public/pyris/pipelines/tutor-chat/runs/" + jobId + "/status", new PyrisChatStatusUpdateDTO(result, stages, null, tokens),
HttpStatus.OK, headers)
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
🛑 Comments failed to post (30)
src/main/java/de/tum/cit/aet/artemis/core/domain/LLMRequest.java (1)
3-3: 🛠️ Refactor suggestion
Consider using BigDecimal for cost-related fields.
For the
costPerMillionInputTokenandcostPerMillionOutputTokenfields, consider usingjava.math.BigDecimalinstead offloat. BigDecimal is more suitable for representing currency and performing precise decimal calculations, which is crucial for cost-related data.Here's the suggested change:
- public record LLMRequest(String model, int numInputTokens, float costPerMillionInputToken, int numOutputTokens, float costPerMillionOutputToken, String pipelineId) { + public record LLMRequest(String model, int numInputTokens, BigDecimal costPerMillionInputToken, int numOutputTokens, BigDecimal costPerMillionOutputToken, String pipelineId) {Don't forget to add the necessary import:
import java.math.BigDecimal;src/main/java/de/tum/cit/aet/artemis/athena/dto/ResponseMetaDTO.java (1)
15-17: 🛠️ Refactor suggestion
Consider using primitives and BigDecimal for improved efficiency and precision.
The
TotalUsagerecord is well-structured, but consider the following suggestions:
- If null values are not needed for these fields, using primitive types (int and float) instead of wrapper classes could improve performance.
- For the
costfield, consider usingBigDecimalinstead ofFloatto ensure precise financial calculations.Here's a potential refactoring:
public record TotalUsage(int numInputTokens, int numOutputTokens, int numTotalTokens, BigDecimal cost) { }Make sure to import
java.math.BigDecimalif you decide to use it.src/test/java/de/tum/cit/aet/artemis/iris/IrisChatWebsocketTest.java (1)
56-56: 💡 Codebase verification
Missing verification of the new parameter in IrisChatTokenTrackingIntegrationTest.java.
The integration test does not verify the new
List.of()parameter in theIrisChatWebsocketDTOconstructor, which may lead to incomplete test coverage.🔗 Analysis chain
LGTM! Consider additional integration tests.
The change correctly updates the
IrisChatWebsocketDTOconstructor call to include the new parameter. This aligns with the coding guidelines and maintains the test's integrity.As per the retrieved learning, consider implementing additional tests for the new field in a separate integration test class, such as
IrisChatTokenTrackingIntegrationTest. This would ensure comprehensive coverage of the new functionality.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Check if an integration test for the new field exists fd -e java -p 'IrisChatTokenTrackingIntegrationTest'Length of output: 136
Script:
#!/bin/bash # Check for usage of List.of() in IrisChatTokenTrackingIntegrationTest.java grep -E 'List\.of\(\)' src/test/java/de/tum/cit/aet/artemis/iris/IrisChatTokenTrackingIntegrationTest.javaLength of output: 108
Script:
#!/bin/bash # Search for IrisChatWebsocketDTO constructor calls with multiple parameters grep -E 'new IrisChatWebsocketDTO\s*\(.*\)' src/test/java/de/tum/cit/aet/artemis/iris/IrisChatTokenTrackingIntegrationTest.javaLength of output: 129
src/main/java/de/tum/cit/aet/artemis/core/domain/LLMTokenUsageRequest.java (4)
28-30:
⚠️ Potential issueUse
BigDecimalinstead offloatfor monetary valuesUsing
floatfor monetary values can lead to precision and rounding errors due to floating-point representation. It is recommended to useBigDecimalfor precise monetary calculations.Apply this diff to change the data types:
@Column(name = "cost_per_million_input_tokens") - private float costPerMillionInputTokens; + private BigDecimal costPerMillionInputTokens; @Column(name = "cost_per_million_output_tokens") - private float costPerMillionOutputTokens; + private BigDecimal costPerMillionOutputTokens;Remember to update the getter and setter methods accordingly to reflect the new data type.
Also applies to: 34-36
68-70:
⚠️ Potential issueInconsistent parameter name in setter
setCostPerMillionOutputTokensThe parameter name
costPerMillionOutputToken(singular) does not match the field namecostPerMillionOutputTokens(plural). For consistency and clarity, the parameter name should be updated tocostPerMillionOutputTokens.Apply this diff to fix the parameter name:
public void setCostPerMillionOutputTokens(float costPerMillionOutputToken) { - this.costPerMillionOutputTokens = costPerMillionOutputToken; + this.costPerMillionOutputTokens = costPerMillionOutputTokens; } +// Updated method signature +public void setCostPerMillionOutputTokens(float costPerMillionOutputTokens) { + this.costPerMillionOutputTokens = costPerMillionOutputTokens; +}Committable suggestion was skipped due to low confidence.
60-62:
⚠️ Potential issueInconsistent parameter name in setter
setCostPerMillionInputTokensThe parameter name
costPerMillionInputToken(singular) does not match the field namecostPerMillionInputTokens(plural). For consistency and clarity, the parameter name should be updated tocostPerMillionInputTokens.Apply this diff to fix the parameter name:
public void setCostPerMillionInputTokens(float costPerMillionInputToken) { - this.costPerMillionInputTokens = costPerMillionInputToken; + this.costPerMillionInputTokens = costPerMillionInputTokens; } +// Updated method signature +public void setCostPerMillionInputTokens(float costPerMillionInputTokens) { + this.costPerMillionInputTokens = costPerMillionInputTokens; +}Committable suggestion was skipped due to low confidence.
19-94: 🛠️ Refactor suggestion
Consider adding constructor injection for immutability
To enhance the immutability of the class and adhere to the dependency injection principle, consider adding constructor parameters for required fields. This could make the class instances immutable once created.
Here's how you might modify the class:
public class LLMTokenUsageRequest extends DomainObject { + // Constructor with required fields + public LLMTokenUsageRequest(String model, String servicePipelineId, int numInputTokens, BigDecimal costPerMillionInputTokens, int numOutputTokens, BigDecimal costPerMillionOutputTokens, LLMTokenUsageTrace trace) { + this.model = model; + this.servicePipelineId = servicePipelineId; + this.numInputTokens = numInputTokens; + this.costPerMillionInputTokens = costPerMillionInputTokens; + this.numOutputTokens = numOutputTokens; + this.costPerMillionOutputTokens = costPerMillionOutputTokens; + this.trace = trace; + } // Existing fields and methods... }Committable suggestion was skipped due to low confidence.
src/main/java/de/tum/cit/aet/artemis/core/domain/LLMTokenUsageTrace.java (3)
41-41:
⚠️ Potential issueFix inconsistent data types for 'userId' field and its getter/setter.
The
userIdfield is declared asLong, but the getter and setter use the primitivelong. This inconsistency may lead toNullPointerExceptionifuserIdisnull.Apply this diff to align the data types:
Option 1: If
userIdshould never benull, change the field to use the primitivelong:- private Long userId; + private long userId; public long getUserId() { return userId; } public void setUserId(long userId) { this.userId = userId; }Option 2: If
userIdcan benull, update the getter and setter to useLong:public long getUserId() { - return userId; + return userId; } public void setUserId(long userId) { - this.userId = userId; + this.userId = userId; }Ensure consistency between the field and its accessor methods.
Also applies to: 77-83
43-44: 🛠️ Refactor suggestion
Consider using
@CreationTimestampfor the 'time' field.Initializing the
timefield withZonedDateTime.now()may lead to unexpected results if the entity is not immediately persisted or if the object is reused. Using@CreationTimestampensures that the timestamp is automatically set when the entity is saved to the database.Apply this diff:
+ @CreationTimestamp @Column(name = "time") - private ZonedDateTime time = ZonedDateTime.now(); + private ZonedDateTime time;Remember to import
org.hibernate.annotations.CreationTimestamp.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.@CreationTimestamp @Column(name = "time") private ZonedDateTime time;
93-99:
⚠️ Potential issueEnsure consistent method naming for 'llmRequests' getters and setters.
The getter and setter methods for
llmRequestshave inconsistent naming conventions. The getter isgetLLMRequests(), while the setter issetLlmRequests(). This inconsistency may cause issues with frameworks that rely on JavaBean naming conventions.Apply this diff to standardize the method names:
Option 1: Use
getLlmRequests()andsetLlmRequests()to match the field name:- public Set<LLMTokenUsageRequest> getLLMRequests() { + public Set<LLMTokenUsageRequest> getLlmRequests() { return llmRequests; } public void setLlmRequests(Set<LLMTokenUsageRequest> llmRequests) { this.llmRequests = llmRequests; }Option 2: Use
getLLMRequests()andsetLLMRequests()to capitalizeLLM:public Set<LLMTokenUsageRequest> getLLMRequests() { return llmRequests; } - public void setLlmRequests(Set<LLMTokenUsageRequest> llmRequests) { + public void setLLMRequests(Set<LLMTokenUsageRequest> llmRequests) { this.llmRequests = llmRequests; }Choose one option to maintain consistency and adhere to naming conventions.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.public Set<LLMTokenUsageRequest> getLLMRequests() { return llmRequests; } public void setLLMRequests(Set<LLMTokenUsageRequest> llmRequests) { this.llmRequests = llmRequests; }src/main/java/de/tum/cit/aet/artemis/iris/service/websocket/IrisChatWebsocketService.java (4)
45-45: 🛠️ Refactor suggestion
Avoid passing multiple nulls to the constructor
Passing multiple
nullvalues toIrisChatWebsocketDTOreduces code readability and increases the risk of errors. Consider implementing the Builder pattern or providing overloaded constructors to handle optional parameters more gracefully.
56-56: 🛠️ Refactor suggestion
Eliminate method calls with multiple null arguments
Invoking
sendStatusUpdatewithnullvalues forsuggestionsandtokensdecreases code clarity. Provide overloaded methods or use optional parameters to enhance readability and maintainability.
67-67: 🛠️ Refactor suggestion
Reduce parameter count in
sendStatusUpdatemethodThe
sendStatusUpdatemethod now has four parameters, increasing complexity and the potential for errors. Consider encapsulating related parameters into a data transfer object (DTO) or using a builder to simplify the method signature.
71-71:
⚠️ Potential issueAvoid passing null to constructors
Passing
nullas an argument when constructingIrisChatWebsocketDTOcan lead toNullPointerExceptions and reduces code clarity. Refactor the constructor or use patterns that handle optional fields more effectively.Apply this suggested refactor:
- var payload = new IrisChatWebsocketDTO(null, rateLimitInfo, stages, suggestions, tokens); + var payload = IrisChatWebsocketDTO.builder() + .rateLimitInfo(rateLimitInfo) + .stages(stages) + .suggestions(suggestions) + .tokens(tokens) + .build();Committable suggestion was skipped due to low confidence.
src/main/java/de/tum/cit/aet/artemis/iris/service/session/AbstractIrisChatSessionService.java (3)
34-34:
⚠️ Potential issuePotential concurrency issue with non-thread-safe
HashMapThe
tracesHashMap is accessed in a concurrent context but is not thread-safe. If multiple threads interact with this map, it could lead to data inconsistency or race conditions. Consider using aConcurrentHashMapto ensure thread safety.
73-116: 🛠️ Refactor suggestion
Consider splitting
handleStatusUpdateinto smaller methods for clarityThe
handleStatusUpdatemethod handles multiple responsibilities, such as message processing, token usage tracking, and updating suggestions. To adhere to the single responsibility principle and improve readability, consider refactoring this method by extracting parts of the logic into separate private methods.🧰 Tools
🪛 ast-grep
[warning] 77-77: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: (new IrisTextMessageContent(statusUpdate.result()))
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 77-77: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: (new IrisTextMessageContent(statusUpdate.result()))
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
77-78:
⚠️ Potential issueEnsure sanitization of
statusUpdate.result()to prevent security vulnerabilitiesThe
statusUpdate.result()may contain user-generated or external content. Before creating a newIrisTextMessageContentwith this content, ensure that it is properly sanitized or escaped to prevent security issues like cross-site scripting (XSS) attacks.🧰 Tools
🪛 ast-grep
[warning] 77-77: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: (new IrisTextMessageContent(statusUpdate.result()))
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 77-77: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: (new IrisTextMessageContent(statusUpdate.result()))
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfigurationsrc/main/java/de/tum/cit/aet/artemis/core/service/LLMTokenUsageService.java (4)
49-49: 🛠️ Refactor suggestion
Consider making the
saveLLMTokenUsagemethod asynchronous as per the TODO comment.The TODO comment suggests that this method should ideally be executed asynchronously to prevent blocking the main thread, especially if the database operation is time-consuming. Implementing it asynchronously can improve the application's performance and responsiveness. Consider using Spring's
@Asyncannotation or other asynchronous execution mechanisms.Would you like assistance in implementing asynchronous processing for this method?
77-77: 🛠️ Refactor suggestion
Consider making the
appendRequestsToTracemethod asynchronous as per the TODO comment.Similarly, the TODO comment indicates that this method should be executed asynchronously. Making this method asynchronous can enhance performance by freeing up resources and allowing other processes to run concurrently.
Would you like assistance in making this method execute asynchronously?
79-80:
⚠️ Potential issueAvoid using
peekfor side-effects; usemapinstead.As with the previous instance, using
peekfor side-effects is not recommended. Refactor to usemapto modify and transform the elements appropriately.Consider refactoring the code as follows:
- var requestSet = requests.stream().map(LLMTokenUsageService::convertLLMRequestToLLMTokenUsageRequest) - .peek(llmTokenUsageRequest -> llmTokenUsageRequest.setTrace(trace)) - .collect(Collectors.toSet()); + var requestSet = requests.stream() + .map(llmRequest -> { + LLMTokenUsageRequest usageRequest = LLMTokenUsageService.convertLLMRequestToLLMTokenUsageRequest(llmRequest); + usageRequest.setTrace(trace); + return usageRequest; + }) + .collect(Collectors.toSet());📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.var requestSet = requests.stream() .map(llmRequest -> { LLMTokenUsageRequest usageRequest = LLMTokenUsageService.convertLLMRequestToLLMTokenUsageRequest(llmRequest); usageRequest.setTrace(trace); return usageRequest; }) .collect(Collectors.toSet());
60-61:
⚠️ Potential issueAvoid using
peekfor side-effects; usemapinstead.The
peekmethod is intended for debugging and should not be used to perform side-effects like modifying elements. Usingmapis more appropriate when you need to transform elements, especially when the transformation includes side-effects.Consider refactoring the code as follows:
- llmTokenUsageTrace.setLlmRequests(llmRequests.stream().map(LLMTokenUsageService::convertLLMRequestToLLMTokenUsageRequest) - .peek(llmTokenUsageRequest -> llmTokenUsageRequest.setTrace(llmTokenUsageTrace)).collect(Collectors.toSet())); + llmTokenUsageTrace.setLlmRequests(llmRequests.stream() + .map(llmRequest -> { + LLMTokenUsageRequest usageRequest = LLMTokenUsageService.convertLLMRequestToLLMTokenUsageRequest(llmRequest); + usageRequest.setTrace(llmTokenUsageTrace); + return usageRequest; + }) + .collect(Collectors.toSet()));📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.llmTokenUsageTrace.setLlmRequests(llmRequests.stream() .map(llmRequest -> { LLMTokenUsageRequest usageRequest = LLMTokenUsageService.convertLLMRequestToLLMTokenUsageRequest(llmRequest); usageRequest.setTrace(llmTokenUsageTrace); return usageRequest; }) .collect(Collectors.toSet()));src/main/java/de/tum/cit/aet/artemis/iris/service/session/IrisExerciseChatSessionService.java (2)
60-65: 🛠️ Refactor suggestion
Constructor has too many parameters
The constructor now has twelve parameters, which can make the code harder to maintain and test. Consider refactoring to reduce the number of parameters, possibly by grouping related services into composite objects or utilizing a dependency injection framework to manage dependencies more efficiently.
155-158:
⚠️ Potential issueEnsure safe access to
exerciseto prevent potentialNullPointerExceptionTo prevent a possible
NullPointerException, ensure thatexercise.getCourseViaExerciseGroupOrCourseMember()does not returnnullbefore callinggetId(). Consider adding a null check or handling potential null values appropriately.src/main/java/de/tum/cit/aet/artemis/athena/service/AthenaFeedbackSuggestionsService.java (2)
169-181:
⚠️ Potential issuePotential
ClassCastExceptionwhen casting toStudentParticipationAt line 174, the code casts
submission.getParticipation()toStudentParticipationwithout an instance check. This could lead to aClassCastExceptionif the participation is not of typeStudentParticipation, such as when it's aTeacherParticipation.Consider verifying the instance before casting:
Participation participation = submission.getParticipation(); Long userId = null; if (participation instanceof StudentParticipation studentParticipation) { userId = studentParticipation.getStudent().map(User::getId).orElse(null); }This ensures safe casting and prevents potential runtime exceptions.
66-66:
⚠️ Potential issueTypo in parameter documentation
There's a typo in the documentation for
athenaDTOConverterService: "exrcises" should be "exercises".Apply this diff to fix the typo:
- * @param athenaDTOConverterService Service to convert exrcises and submissions to DTOs + * @param athenaDTOConverterService Service to convert exercises and submissions to DTOs📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.* @param athenaDTOConverterService Service to convert exercises and submissions to DTOssrc/test/java/de/tum/cit/aet/artemis/iris/IrisChatTokenTrackingIntegrationTest.java (5)
155-157: 🛠️ Refactor suggestion
Improve test method name for specificity
The method
testTokenTrackingSavedExerciseChatcan be renamed to more specifically describe the test scenario. A name likeshouldSaveTokenUsageWhenExerciseChatMessageIsStoredwould provide clearer insight into the test's intention.
171-173: 🛠️ Refactor suggestion
Make test method name more descriptive
To enhance clarity, consider renaming
testTokenTrackingExerciseChatWithPipelineFailtoshouldTrackTokensWhenPipelineFailsDuringExerciseChatto accurately reflect the test scenario and expected outcome.
79-119: 🛠️ Refactor suggestion
Avoid direct database access in tests
The
initTestCasemethod directly interacts with repositories to manipulate the database, which is discouraged by the guidelineavoid_db_access: truefor test files. Direct database access can slow down tests and make them less reliable. Consider using mock repositories or test utilities that simulate database operations to improve test performance and maintainability.
124-135: 🛠️ Refactor suggestion
Refactor duplicated code in test methods
The setup and execution logic in
testTokenTrackingHandledExerciseChatandtestTokenTrackingExerciseChatWithPipelineFailare largely similar. To adhere to the DRY (Don't Repeat Yourself) principle and improve maintainability, consider extracting common code into helper methods or setup functions. This refactoring will reduce duplication and make future updates easier.Also applies to: 174-185
🧰 Tools
🪛 ast-grep
[warning] 127-127: Detected a cookie where the
Secureflag is either missing or disabled. TheSecurecookie flag instructs the browser to forbid sending the cookie over an insecure HTTP request. Set theSecureflag totrueso the cookie will only be sent over HTTPS.
Context: (new PyrisStageDTO("DoneTest", 10, PyrisStageState.DONE, "Done"))
Note: [CWE-614]: Sensitive Cookie in HTTPS Session Without 'Secure' Attribute [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
[warning] 127-127: Detected a cookie where the
HttpOnlyflag is either missing or disabled. TheHttpOnlycookie flag instructs the browser to forbid client-side JavaScript to read the cookie. If JavaScript interaction is required, you can ignore this finding. However, set theHttpOnlyflag to true` in all other cases.
Context: (new PyrisStageDTO("DoneTest", 10, PyrisStageState.DONE, "Done"))
Note: [CWE-1004]: Sensitive Cookie Without 'HttpOnly' Flag [OWASP A05:2021]: Security Misconfiguration [REFERENCES]
- https://owasp.org/Top10/A05_2021-Security_Misconfiguration
121-123: 🛠️ Refactor suggestion
Enhance test method naming for clarity
The test method
testTokenTrackingHandledExerciseChatcould be renamed to more precisely reflect its purpose. Consider renaming it toshouldTrackTokensWhenHandlingExerciseChatfor better readability and to clearly indicate the expected behavior being tested.
bassner
merged commit cc127af
into
feature/track-usage-of-iris-requests
Checklist
General
Server
Motivation
We want to track the token usage for Athena.
Description
Stores token usage of Athena from the response meta field for feedback generation requests.
Follow-up Issues
costPerMillionInputToken,costPerMillionOutputTokenSteps for Testing
You'd need to request feedback using Athena and check the database tables if the calls are correctly stored.
This can be done locally.
Testserver States
Note
These badges show the state of the test servers.
Green = Currently available, Red = Currently locked
Click on the badges to get to the test servers.
Review Progress
Performance Review
Code Review
Manual Tests
Exam Mode Test
Performance Tests
Test Coverage
Screenshots
Summary by CodeRabbit
New Features
ResponseMetaDTOto encapsulate metadata related to responses from the Athena service.LLMRequestandLLMServiceTypeto manage requests and service types for large language models.LLMTokenUsageServicefor tracking token usage within the system.IrisChatWebsocketDTOto include token information for improved WebSocket communication.Bug Fixes
Tests
Documentation