Track token usage of iris requests #165
Conversation
Fix last lint error Fix last lint errors
…ests' into track-token-usage
WalkthroughThis pull request introduces the Changes
Possibly related PRs
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
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Outside diff range and nitpick comments (21)
app/domain/data/token_usage_dto.py (1)
6-10: LGTM: Well-structured DTO with a minor suggestion.The
TokenUsageDTOclass is well-defined and aligns with the PR objectives:
- Inherits from
pydantic.BaseModelfor robust data validation.- Clear and descriptive attribute names with appropriate type annotations.
- The
pipelineattribute usingPipelineEnumsupports tracking token usage across different pipelines.Consider adding a docstring to the class to provide more context and usage information. For example:
class TokenUsageDTO(BaseModel): """ Data Transfer Object for token usage information. Attributes: model_info (str): Information about the model used. num_input_tokens (int): Number of input tokens processed. num_output_tokens (int): Number of output tokens generated. pipeline (PipelineEnum): The pipeline associated with this token usage. """ # ... (rest of the class definition)app/llm/external/PipelineEnum.py (1)
5-14: LGTM: Well-structured enum members with a minor suggestion.The enum members are well-defined and consistently named. They cover various aspects of the IRIS system, which aligns with the PR objectives. The inclusion of a
NOT_SETmember is a good practice for handling undefined cases.Consider adding a brief docstring to the
PipelineEnumclass to describe its purpose and the meaning of each member. This would enhance the code's self-documentation. For example:class PipelineEnum(str, Enum): """ Enumeration of pipeline types in the IRIS system. Each member represents a specific pipeline or processing stage used for tracking token usage and other metrics. """ # ... enum members ...app/llm/external/LLMTokenCount.py (1)
23-24: LGTM with a minor suggestion: Consider enhancing the__str__method.The
__str__method provides a clear and informative string representation of the object, which is great. However, to improve consistency and readability, consider using f-string's formatting options for better alignment. Here's a suggested improvement:def __str__(self): return (f"{self.model_info}: " f"input_tokens={self.num_input_tokens:<6} " f"output_tokens={self.num_output_tokens:<6} " f"pipeline={self.pipeline.name}")This format aligns the values and uses the enum's name for better readability. It also separates the information more clearly, which can be helpful when logging or debugging.
app/domain/pyris_message.py (1)
19-22: LGTM! Consider reordering fields for better readability.The new fields
num_input_tokens,num_output_tokens, andmodel_infoare well-implemented and align with the PR objectives for tracking token usage. The use of aliases and default values is a good practice for serialization and backward compatibility.Consider reordering the fields to group related attributes together, improving readability. For example:
class PyrisMessage(BaseModel): model_config = ConfigDict(populate_by_name=True) sent_at: datetime | None = Field(alias="sentAt", default=None) sender: IrisMessageRole contents: List[MessageContentDTO] = [] num_input_tokens: int = Field(alias="numInputTokens", default=0) num_output_tokens: int = Field(alias="numOutputTokens", default=0) model_info: str = Field(alias="modelInfo", default="")This groups the message-related fields together, followed by the token and model information fields.
app/pipeline/pipeline.py (1)
11-11: Consider initializingtokensin__init__method.The addition of the
tokensattribute aligns with the PR objectives for tracking token usage. However, there are a couple of points to consider:
- The attribute is defined as a class attribute rather than an instance attribute. This might lead to unexpected behavior if shared across instances.
- The
__init__method hasn't been updated to initialize this attribute.Consider modifying the
__init__method to initializetokensas an instance attribute:class Pipeline(metaclass=ABCMeta): """Abstract class for all pipelines""" implementation_id: str - tokens: List[LLMTokenCount] - def __init__(self, implementation_id=None, **kwargs): + def __init__(self, implementation_id=None, tokens: List[LLMTokenCount] = None, **kwargs): self.implementation_id = implementation_id + self.tokens = tokens if tokens is not None else []This change ensures that each instance has its own
tokenslist and allows for optional initialization.app/llm/langchain/iris_langchain_chat_model.py (1)
Line range hint
19-20: Consider updating the class docstring.The class has been enhanced with token tracking functionality. It would be beneficial to update the class docstring to reflect this new capability. This will help maintain clear and up-to-date documentation.
Here's a suggested addition to the docstring:
class IrisLangchainChatModel(BaseChatModel): """Custom langchain chat model for our own request handler. This class now includes token tracking functionality, storing token usage information in the `tokens` attribute after each generation. """app/pipeline/shared/reranker_pipeline.py (1)
60-60: Consider adding a type hint forself.tokensWhile the initialization of
self.tokensis correct and aligns with the PR objectives, adding a type hint would improve code clarity and maintainability.Consider updating the line as follows:
self.tokens: List[LLMTokenCount] = []Don't forget to import
Listfromtypingif not already imported.app/web/status/status_update.py (2)
Line range hint
101-113: LGTM: Enhanceddonemethod with token trackingThe changes to the
donemethod successfully incorporate token tracking and provide more control over the next stage. The implementation looks correct and aligns with the PR objectives.Consider adding a brief comment explaining the purpose of the new parameters, especially
tokens, to improve code documentation.
Issue: Some
StatusCallback.done()calls are missing thetokensargumentSeveral instances of
StatusCallback.done()do not include thetokensparameter, which may lead to inconsistent token tracking.Affected files:
app/pipeline/lecture_ingestion_pipeline.py
- Line 123
- Line 142
- Line 145
app/pipeline/chat/exercise_chat_pipeline.py
- Line 102
- Line 119
- Line 245
- Line 270
app/pipeline/chat/course_chat_pipeline.py
- Line 422
Please update these calls to include the
tokensparameter as necessary to ensure consistent token tracking across the codebase.🔗 Analysis chain
Line range hint
22-113: Overall: Successfully implemented token trackingThe changes successfully add token tracking functionality to the
StatusCallbackclass, aligning with the PR objectives. The implementation is backward-compatible and minimally invasive.To ensure full utilization of this new feature:
- Verify that other parts of the codebase that use the
StatusCallbackclass are updated to pass token information when available.- Consider adding unit tests to cover the new token tracking functionality.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for usages of StatusCallback.done() method that might need updating # Test: Search for calls to StatusCallback.done() or its subclasses echo "Searching for StatusCallback.done() calls:" rg --type python -n "\.done\(" --glob '!app/web/status/status_update.py' echo "Searching for potential subclasses of StatusCallback:" rg --type python -n "class.*\(.*StatusCallback.*\):"Length of output: 403
Script:
#!/bin/bash # Description: Check for usages of StatusCallback.done() method that might need updating # Test: Search for calls to StatusCallback.done() or its subclasses echo "Searching for StatusCallback.done() calls:" rg "\.done\(" --glob '*.py' -n --ignore-file '!app/web/status/status_update.py' echo "Searching for potential subclasses of StatusCallback:" rg "class\s+\w+\s*\(.*StatusCallback.*\):" --glob '*.py' -nLength of output: 1841
Script:
#!/bin/bash # Description: Verify that all StatusCallback.done() method calls include the 'tokens' argument # Search for calls to StatusCallback.done() excluding the status_update.py file echo "Searching for StatusCallback.done() calls excluding app/web/status/status_update.py:" rg "\.done\(" --glob '*.py' --glob '!app/web/status/status_update.py' -n # Extract the arguments of each done() call to check for 'tokens' echo "Extracting arguments of StatusCallback.done() calls to verify 'tokens' parameter inclusion:" rg "\.done\(" --glob '*.py' --glob '!app/web/status/status_update.py' -n -o "done\(([^)]*)\)"Length of output: 1667
app/pipeline/chat/course_chat_pipeline.py (1)
420-422: LGTM: Comprehensive token tracking and reportingThe implementation successfully:
- Includes tokens from the citation pipeline in the total count.
- Sends token usage data to Artemis as part of the status update.
These changes align well with the PR objectives. However, consider adding a comment explaining the purpose of extending
self.tokenswithself.citation_pipeline.tokensfor better code readability.Consider adding a comment above line 420 to explain the purpose of including citation pipeline tokens:
# Include tokens from citation pipeline in the total count self.tokens.extend(self.citation_pipeline.tokens)app/pipeline/shared/citation_pipeline.py (3)
Line range hint
2-2: Incorrect import ofloggerThe import statement
from asyncio.log import loggeris incorrect because theasynciomodule does not have alogsubmodule that provideslogger. Instead, you should use Python's built-inloggingmodule. Replace the import statement with:-from asyncio.log import logger +import logging +logger = logging.getLogger(__name__)
92-92: Use logging instead ofThe use of
print(response)is not recommended in production code. It is better to use a logger for controlled and configurable output. Replace the- print(response) + logger.info(f"Citation pipeline response: {response}")
Line range hint
100-102: Improve exception handling and loggingIn the exception block, the statement
logger.error("citation pipeline failed", e)may not correctly log the exception details. Uselogger.exceptionto include the traceback automatically:- except Exception as e: - logger.error("citation pipeline failed", e) + except Exception: + logger.exception("Citation pipeline failed")This will provide more detailed information about the exception, aiding in debugging.
app/llm/external/ollama.py (1)
Line range hint
60-73: Ensure 'message' supports attribute access instead of subscriptable accessIn the
convert_to_iris_messagefunction, you're accessingmessageusingmessage["content"]andmessage["role"]. Ifmessageis an instance of theMessageclass, it may not support subscriptable access (__getitem__). Instead, use attribute access likemessage.contentandmessage.role.Apply this diff to correct the attribute access:
def convert_to_iris_message( message: Message, num_input_tokens: int, num_output_tokens: int, model: str ) -> PyrisMessage: """ Convert a Message to a PyrisMessage """ - contents = [TextMessageContentDTO(text_content=message["content"])] + contents = [TextMessageContentDTO(text_content=message.content)] return PyrisMessage( - sender=map_str_to_role(message["role"]), + sender=map_str_to_role(message.role), contents=contents, send_at=datetime.now(), num_input_tokens=num_input_tokens, num_output_tokens=num_output_tokens, model_info=model, )app/pipeline/chat/code_feedback_pipeline.py (1)
45-45: Consider initializingtokensin the__init__methodWhile declaring
tokens: LLMTokenCountat the class level defines its type, initializingself.tokensin the__init__method ensures that the attribute is instantiated when an object is created.You may apply this change:
def __init__(self, callback: Optional[StatusCallback] = None): super().__init__(implementation_id="code_feedback_pipeline_reference_impl") # existing code... + self.tokens = LLMTokenCount()app/llm/external/openai_chat.py (3)
71-72: Consider defaulting token counts to zero instead of -1When
usagedoes not includeprompt_tokensorcompletion_tokens, the default value is set to-1. This might lead to confusion when interpreting token counts, as negative values may not be expected in token usage metrics.Consider setting the default values to zero to accurately represent the absence of token usage:
-num_input_tokens = getattr(usage, "prompt_tokens", -1) -num_output_tokens = getattr(usage, "completion_tokens", -1) +num_input_tokens = getattr(usage, "prompt_tokens", 0) +num_output_tokens = getattr(usage, "completion_tokens", 0)
Line range hint
105-105: Uselogginginstead ofThe use of
loggingmodule provides better control over logging levels and formats.Replace the
logging.infocall to align with the logging practices used elsewhere in the code:-print("Sending messages to OpenAI", messages) +logging.info("Sending messages to OpenAI: %s", messages)
Line range hint
125-125: Correct the error message in exception handlingThe error message states "Failed to interpret image after several attempts," which is not relevant to the chat functionality. This appears to be a copy-paste error and could cause confusion during debugging.
Update the error message to accurately reflect the operation being performed:
-logging.error("Failed to interpret image after several attempts.") +logging.error("Failed to receive a response from OpenAI after several attempts.")app/pipeline/chat/interaction_suggestion_pipeline.py (1)
Line range hint
57-171: Consider adding unit tests for the new token tracking functionality.To ensure the robustness of the token tracking feature, please add unit tests that verify:
self.tokensis correctly assigned after pipeline execution.self.tokens.pipelineis set to the appropriatePipelineEnumvalue.- The token counts reflect the actual usage during the pipeline run.
app/pipeline/lecture_ingestion_pipeline.py (2)
9-9: Potential naming conflict between imported module and variabletokensYou've imported
tokensfromsipbuild.generator.parser.tokens, and later usedtokensas a variable name (self.tokens). This could lead to confusion or unintended behavior if the importedtokensis used elsewhere in the code. Consider renaming either the import or the variable to avoid naming conflicts.
282-282: Address the TODO: Implement sending token usage data to ArtemisThere's a TODO comment indicating that the implementation for sending token usage data to Artemis is pending. Completing this functionality is crucial for tracking and reporting token usage as intended by the PR objectives.
Would you like assistance in generating the code to send
self.tokensto Artemis? I can help draft the necessary implementation or open a new GitHub issue to track this task.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (21)
- app/domain/data/token_usage_dto.py (1 hunks)
- app/domain/pyris_message.py (1 hunks)
- app/domain/status/status_update_dto.py (1 hunks)
- app/llm/external/LLMTokenCount.py (1 hunks)
- app/llm/external/PipelineEnum.py (1 hunks)
- app/llm/external/ollama.py (3 hunks)
- app/llm/external/openai_chat.py (3 hunks)
- app/llm/langchain/iris_langchain_chat_model.py (3 hunks)
- app/pipeline/chat/code_feedback_pipeline.py (3 hunks)
- app/pipeline/chat/course_chat_pipeline.py (3 hunks)
- app/pipeline/chat/exercise_chat_pipeline.py (8 hunks)
- app/pipeline/chat/interaction_suggestion_pipeline.py (3 hunks)
- app/pipeline/chat/lecture_chat_pipeline.py (3 hunks)
- app/pipeline/competency_extraction_pipeline.py (4 hunks)
- app/pipeline/lecture_ingestion_pipeline.py (4 hunks)
- app/pipeline/pipeline.py (1 hunks)
- app/pipeline/shared/citation_pipeline.py (3 hunks)
- app/pipeline/shared/reranker_pipeline.py (3 hunks)
- app/pipeline/shared/summary_pipeline.py (1 hunks)
- app/retrieval/lecture_retrieval.py (3 hunks)
- app/web/status/status_update.py (3 hunks)
🧰 Additional context used
🔇 Additional comments (47)
app/domain/data/token_usage_dto.py (2)
1-3: LGTM: Imports are appropriate and concise.The imports are well-chosen for the purpose of this DTO:
pydantic.BaseModelprovides robust data validation and serialization capabilities.PipelineEnumis imported from a local module, which is good for maintaining a clear project structure.
1-10: Overall assessment: Well-implemented DTO for token usage tracking.This new file introduces a clean and well-structured
TokenUsageDTOthat aligns perfectly with the PR objectives of tracking token usage across different pipelines. The use of Pydantic for the DTO provides robust data validation and serialization capabilities, which is a good practice.The only minor suggestion is to add a docstring to the class for improved documentation.
Great job on implementing this crucial component for enhancing the tracking of iris request usage!
app/domain/status/status_update_dto.py (3)
5-5: LGTM: New import for TokenUsageDTOThe import of
TokenUsageDTOis correctly placed and necessary for the newtokensattribute in theStatusUpdateDTOclass. This change aligns with the PR objective of tracking token usage.
11-11: LGTM: New tokens attribute added to StatusUpdateDTOThe new
tokensattribute in theStatusUpdateDTOclass is well-implemented:
- It uses a clear and descriptive name.
- The type annotation (List[TokenUsageDTO]) is correct and matches the imported type.
- The default value of an empty list is a good practice for optional attributes.
This addition aligns with the PR objective of tracking token usage and sending it to Artemis as part of the status update.
Line range hint
1-11: Summary: Changes effectively support token usage trackingThe modifications to
StatusUpdateDTOare minimal, focused, and directly support the PR objectives:
- The new
tokensattribute allows for the inclusion of token usage data in status updates.- This change enables Pyris to track and report token usage to Artemis, as mentioned in the PR objectives.
The implementation is clean, follows good practices, and integrates well with the existing code structure. These changes lay the groundwork for improved monitoring and management of token usage across different pipelines.
app/llm/external/PipelineEnum.py (3)
1-1: LGTM: Import statement is correct.The import of
Enumfrom theenummodule is appropriate for defining an enumeration class.
4-4: LGTM: Well-defined enum class.The
PipelineEnumclass is well-defined, inheriting from bothstrandEnum. This approach creates a string-based enum, which allows for easy string comparison and serialization while maintaining the benefits of an enumeration.
1-14: Summary: PipelineEnum aligns well with PR objectives.The introduction of
PipelineEnumis a solid foundation for tracking token usage across different pipeline types in the IRIS system. This enum class directly supports the PR's objective of enhancing the tracking of usage related to iris requests.The implementation is clean, follows best practices for string-based enums in Python, and provides a comprehensive list of pipeline types. This will likely be used in conjunction with the
LLMTokenCounttype mentioned in the PR description to facilitate detailed token usage tracking.Great job on laying the groundwork for improved usage monitoring!
app/llm/external/LLMTokenCount.py (4)
1-1: LGTM: Import statement is correct and necessary.The import of
PipelineEnumis properly done and is required for the class implementation.
4-9: LGTM: Class definition and attributes are well-structured.The
LLMTokenCountclass is well-defined with clear and appropriate attributes. The use of type hints enhances code readability and aligns with Python best practices. The attributes effectively capture the necessary information for token usage tracking, which is in line with the PR objectives.
11-21: LGTM: Constructor is well-implemented.The
__init__method is correctly implemented, initializing all attributes with the provided values. The use of type hints in the method signature is consistent with the attribute definitions, which is good practice.
1-24: Overall, excellent implementation ofLLMTokenCountclass.The
LLMTokenCountclass is well-designed and implemented, providing a solid foundation for token usage tracking across different pipelines. It aligns perfectly with the PR objectives and integrates well with the existing system (e.g., usingPipelineEnum). The use of type hints, clear attribute definitions, and informative string representation contribute to code readability and maintainability.A few key strengths:
- Clear and appropriate attribute definitions
- Well-structured
__init__method- Informative
__str__methodThe class forms a crucial part of the token usage tracking system, which will be valuable for monitoring and managing resources across different pipelines.
app/domain/pyris_message.py (1)
19-21: Verify usage of new token and model info fields.The new fields
num_input_tokens,num_output_tokens, andmodel_infohave been added to thePyrisMessageclass. To ensure these changes are fully integrated:
- Verify that these fields are being populated correctly where
PyrisMessageobjects are created or updated.- Check if any existing code needs to be updated to handle these new fields.
- Ensure that any serialization/deserialization logic elsewhere in the codebase is updated to include these new fields.
Run the following script to check for usage of the
PyrisMessageclass and the new fields:This will help identify areas of the codebase that might need updates to fully utilize these new fields.
✅ Verification successful
New token and model info fields are properly utilized across the codebase.
All usages of
num_input_tokens,num_output_tokens, andmodel_infoin thePyrisMessageclass have been verified and correctly implemented in the relevant files.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check usage of PyrisMessage and new fields # Search for PyrisMessage usage echo "PyrisMessage usage:" rg --type python -A 5 "PyrisMessage" # Search for the new field names (both snake_case and camelCase) echo "\nNew field usage:" rg --type python -A 5 -e "num_input_tokens" -e "numInputTokens" -e "num_output_tokens" -e "numOutputTokens" -e "model_info" -e "modelInfo"Length of output: 323
Script:
#!/bin/bash # Description: Check usage of PyrisMessage and new fields by targeting .py files # Search for PyrisMessage usage echo "PyrisMessage usage:" rg -g "*.py" -A 5 "PyrisMessage" # Search for the new field names (both snake_case and camelCase) echo "\nNew field usage:" rg -g "*.py" -A 5 -e "num_input_tokens" -e "numInputTokens" -e "num_output_tokens" -e "numOutputTokens" -e "model_info" -e "modelInfo"Length of output: 33101
app/pipeline/pipeline.py (2)
2-2: LGTM: Import ofListis correct and necessary.The addition of
from typing import Listis appropriate for type hinting the newtokensattribute. This follows Python best practices for type annotations.
Line range hint
1-35: Clarify token tracking usage in pipeline execution.The changes to add token tracking capability are consistent with the PR objectives. However, it's not clear from this file how the
tokensattribute will be populated or used during pipeline execution.To understand how token tracking is integrated into the pipeline execution, please run the following verification script:
#!/bin/bash # Search for usage of 'tokens' attribute in pipeline execution rg --type python -g '!app/pipeline/pipeline.py' 'self\.tokens' app/pipelineConsider adding a comment or documentation to explain how subclasses should interact with the
tokensattribute during execution.app/llm/langchain/iris_langchain_chat_model.py (2)
10-11: LGTM: New imports for token counting.The new imports for
LLMTokenCountandPipelineEnumare correctly placed and align with the PR objectives to track token usage.
24-24: LGTM: New class attribute for token tracking.The
tokensattribute of typeLLMTokenCountis correctly added and initialized. This aligns with the PR objectives and follows good practices for type hinting.app/pipeline/shared/summary_pipeline.py (1)
Line range hint
1-72: Review other pipeline classes for consistencyWhile we've addressed the token tracking implementation in the
SummaryPipelineclass, it's important to ensure that these changes are applied consistently across all pipeline classes in the project. This will ensure that the PR objectives of adding token tracking to all Pipelines are fully met.To verify the consistency of implementation across other pipeline classes, you can run the following script:
This script will help identify any inconsistencies in the implementation of token tracking across different pipeline classes.
✅ Verification successful
Run Updated Verification Scripts
To ensure consistent implementation of token tracking across all pipeline classes, please execute the following scripts:
Token Tracking Implementation Verified Across Pipeline Classes
The token tracking functionality has been consistently implemented across all pipeline classes in the project. All relevant classes initialize and utilize
self.tokens, andLLMTokenCountis properly imported where necessary.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for consistent implementation of token tracking across pipeline classes # Search for pipeline classes and their __init__ methods echo "Checking pipeline classes for token tracking implementation:" rg --type python -A 10 "class \w+Pipeline\(Pipeline\):" | rg -A 10 "def __init__" # Search for usage of self.tokens in pipeline classes echo "\nChecking usage of self.tokens in pipeline classes:" rg --type python "self\.tokens" --glob "*pipeline.py" # Search for LLMTokenCount import in pipeline files echo "\nChecking for LLMTokenCount import in pipeline files:" rg --type python "from .+import.+LLMTokenCount" --glob "*pipeline.py"Length of output: 669
Script:
#!/bin/bash # Description: Check for consistent implementation of token tracking across pipeline classes # Search for pipeline classes and their __init__ methods echo "Checking pipeline classes for token tracking implementation:" rg "class \w+Pipeline\(Pipeline\):" --glob "*.py" -A 10 | rg "def __init__" # Search for usage of self.tokens in pipeline classes echo "\nChecking usage of self.tokens in pipeline classes:" rg "self\.tokens" --glob "*pipeline.py" # Search for LLMTokenCount import in pipeline files echo "\nChecking for LLMTokenCount import in pipeline files:" rg "from .+import.+LLMTokenCount" --glob "*pipeline.py"Length of output: 3877
app/pipeline/shared/reranker_pipeline.py (2)
12-12: LGTM: New import for PipelineEnumThe import of
PipelineEnumaligns with the PR objectives and is correctly placed within the import section.
Line range hint
1-115: Overall assessment: Changes align well with PR objectivesThe modifications to the
RerankerPipelineclass successfully implement token usage tracking as outlined in the PR objectives. The changes are well-integrated and maintain the existing functionality of the pipeline.Key points:
- The new import and use of
PipelineEnumare appropriate.- Token tracking is correctly implemented in the
__call__method.- The
self.tokenslist is properly initialized to store token usage data.Consider implementing the suggested improvements:
- Add a type hint for
self.tokens.- Implement error handling for
self.llm.tokensin the__call__method.These enhancements will further improve code clarity and robustness.
app/web/status/status_update.py (1)
22-22: LGTM: New import for LLMTokenCountThe new import for
LLMTokenCountis correctly added and aligns with the PR objectives of tracking token usage.app/pipeline/chat/course_chat_pipeline.py (4)
44-44: LGTM: New import for token trackingThe addition of
PipelineEnumimport is consistent with the PR objectives for tracking token usage across different pipelines.
111-111: LGTM: New attribute for token trackingThe addition of
self.tokensas a list to store token counts is in line with the PR objectives for tracking token usage during pipeline execution.
411-413: LGTM: Token tracking implementationThe implementation of token tracking is well-executed:
- Retrieves the token count from
self.llm.tokens.- Sets the pipeline type using
PipelineEnum.IRIS_CHAT_COURSE_MESSAGE.- Appends the token count to
self.tokenslist.This is consistent with the PR objectives for tracking token usage across different pipelines.
Line range hint
1-489: Overall assessment: Well-implemented token trackingThe changes in this file successfully implement token tracking for the CourseChatPipeline, aligning with the PR objectives. Key points:
- New import and attribute for token tracking.
- Token count collection during pipeline execution.
- Integration with citation pipeline tokens.
- Reporting of token usage to Artemis.
The implementation is clean and consistent. A minor suggestion for improved readability has been made, but overall, the changes are approved and ready for merging.
app/pipeline/shared/citation_pipeline.py (2)
10-10: Verify the import path forPipelineEnumPlease ensure that
PipelineEnumis correctly imported fromapp.llm.external.PipelineEnum. Confirm that this module path is accurate and thatPipelineEnumis accessible from this location.
88-90:⚠️ Potential issueAssigning
pipelineattribute totoken_countmay cause errorsEnsure that
self.llm.tokensreturns an object that allows setting thepipelineattribute. Iftoken_countis an instance of a class without a predefinedpipelineattribute (e.g., an immutable dataclass), this assignment could raise anAttributeError. Verify that theLLMTokenCountclass is designed to have a mutablepipelineattribute.app/pipeline/competency_extraction_pipeline.py (4)
17-18: Imports added for token trackingThe addition of
LLMTokenCountandPipelineEnumimports is appropriate for implementing token usage tracking.
43-43: Initialization ofself.tokensInitializing
self.tokensas an empty list to store token usage data is correct.
111-111: Confirmself.callback.doneaccepts thetokensparameterEnsure that the
donemethod ofself.callbackis updated to accept thetokensparameter to handle token usage data.You can verify this by running:
#!/bin/bash # Description: Verify that 'done' method accepts 'tokens' parameter. # Test: Search for the 'done' method definition in 'CompetencyExtractionCallback'. rg --type python 'def done' -A 5 # Expected: The method signature includes 'tokens' as a parameter.
82-88: Verify thatresponsecontains required token attributesEnsure that the
responseobject fromself.request_handler.chatincludesmodel_info,num_input_tokens, andnum_output_tokensattributes necessary forLLMTokenCount.To confirm, you can run the following script:
app/pipeline/chat/code_feedback_pipeline.py (2)
16-17: Imports for token tracking added correctlyThe imports for
LLMTokenCountandPipelineEnumhave been properly added to support token usage tracking in the pipeline.
147-149: Verify thatLLMTokenCounthas apipelineattributeEnsure that the
LLMTokenCountclass includes apipelineattribute to store the pipeline type. This prevents potentialAttributeErrorwhen settingnum_tokens.pipeline.You can run the following script to verify if the
pipelineattribute exists inLLMTokenCount:✅ Verification successful
Verified:
LLMTokenCountclass includes thepipelineattributeThe
LLMTokenCountclass correctly defines and initializes thepipelineattribute, ensuring that settingnum_tokens.pipelinewill not cause anAttributeError.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that 'LLMTokenCount' class has a 'pipeline' attribute. # Search for the 'LLMTokenCount' class definition with a 'pipeline' attribute ast-grep --lang python --pattern $'class LLMTokenCount: $$$ pipeline: $_ $$$'Length of output: 1457
app/llm/external/openai_chat.py (2)
74-74:⚠️ Potential issueFix parameter name inconsistency in
TextMessageContentDTOThe parameter
textContentis inconsistent with the naming convention used elsewhere, wheretext_contentis preferred. This inconsistency might cause errors if theTextMessageContentDTOclass expectstext_content.Update the parameter name to maintain consistency and prevent potential attribute errors:
-contents=[TextMessageContentDTO(textContent=message.content)], +contents=[TextMessageContentDTO(text_content=message.content)],Likely invalid or redundant comment.
116-118: Ensureresponse.usageandresponse.modelare always availableWhile the code assumes that
response.usageandresponse.modelare present, there could be cases where these attributes are missing. This might lead toAttributeErrorexceptions.Run the following script to verify that all responses from OpenAI include the
usageandmodelattributes:Ensure that the output includes the
usageandmodelfields. If these fields can be absent, consider adding checks in the code to handle such cases gracefully.app/pipeline/chat/interaction_suggestion_pipeline.py (3)
37-38: Import statements added for token tracking.The additions of
LLMTokenCountandPipelineEnumimports are appropriate for token usage tracking.
57-57: Introducedtokensattribute toInteractionSuggestionPipeline.Adding the
tokensattribute to the class enables tracking of token usage, aligning with the goal of monitoring resource utilization.
170-171: Token usage data correctly assigned after pipeline invocation.Assigning
self.tokens = self.llm.tokensand settingself.tokens.pipelineafter invoking the pipeline correctly updates the token usage information.app/pipeline/chat/lecture_chat_pipeline.py (3)
19-19: LGTM!The import statement correctly includes
PipelineEnum, which is necessary for assigning the pipeline attribute to the token usage data.
78-78: LGTM!Initializing
self.tokensas an empty list prepares the pipeline to collect and store token usage data effectively.
125-125: LGTM!Extending
self.tokenswith tokens from thecitation_pipelineeffectively aggregates token usage from both pipelines.app/pipeline/chat/exercise_chat_pipeline.py (5)
37-37: ImportingPipelineEnumThe import statement correctly adds
PipelineEnumfrom the external module. This is necessary for specifying pipeline types in token tracking.
82-82: Initializeself.tokensfor token trackingInitializing
self.tokensas an empty list allows for collecting token usage data throughout the pipeline execution.
103-105: Pass token usage to callbackThe
tokensparameter is appropriately added toself.callback.done()to transmit the collected token usage data.
211-212: Append tokens fromcode_feedback_pipelineTokens from
code_feedback_pipelineare correctly appended toself.tokenswhen available.
233-234: Append tokens fromretrieverTokens from
self.retrieverare correctly appended toself.tokenswhen available.app/retrieval/lecture_retrieval.py (1)
9-9: ImportingLLMTokenCountThe import statement for
LLMTokenCountis appropriate and aligns with the integration of token counting functionality.
| @@ -45,6 +45,7 @@ def __init__(self): | |||
| ) | |||
| # Create the pipeline | |||
| self.pipeline = self.prompt | self.llm | StrOutputParser() | |||
| self.tokens = [] | |||
There was a problem hiding this comment.
Incomplete implementation of token tracking functionality
The addition of self.tokens = [] is a step towards implementing token tracking, but it falls short of the PR objectives. Consider the following points:
- The PR objectives mention adding a
tokensargument to all Pipelines, which is not implemented in the__init__method. - The
self.tokenslist is initialized but never used in the class methods, particularly in the__call__method where token usage should likely be tracked. - There's no implementation of sending token usage data to Artemis as mentioned in the PR objectives.
- The
LLMTokenCounttype mentioned in the PR objectives is not used or imported.
To fully implement the token tracking functionality, consider the following changes:
- Update the
__init__method to accept atokensargument:
def __init__(self, tokens: Optional[List[LLMTokenCount]] = None):
super().__init__(implementation_id="summary_pipeline")
self.tokens = tokens or []- Implement token usage tracking in the
__call__method:
def __call__(self, query: str, **kwargs) -> str:
# ... existing code ...
response: str = self.pipeline.invoke({"text": query})
# Track token usage
token_count = self.llm.get_num_tokens(response) # Assuming this method exists
self.tokens.append(LLMTokenCount(count=token_count, pipeline=PipelineEnum.SUMMARY))
# ... rest of the method ...- Add a method to send token usage to Artemis (this might require additional dependencies):
def send_token_usage_to_artemis(self):
# Implementation to send self.tokens to Artemis
pass- Import necessary types and enums at the top of the file:
from typing import List, Optional
from .token_types import LLMTokenCount, PipelineEnumWould you like me to provide a more detailed implementation or open a GitHub issue to track these changes?
| num_tokens = self.llm.tokens | ||
| num_tokens.pipeline = PipelineEnum.IRIS_RERANKER_PIPELINE | ||
| self.tokens.append(num_tokens) |
There was a problem hiding this comment.
💡 Codebase verification
Potential Missing Error Handling for self.llm.tokens Across Codebase
The recent addition of error handling for self.llm.tokens in app/pipeline/shared/reranker_pipeline.py is a positive step towards ensuring robustness. However, the usage of self.llm.tokens has been identified in multiple other files without similar checks. It's recommended to implement appropriate error handling in all instances where self.llm.tokens is accessed to prevent potential NoneType errors.
Affected Files:
app/retrieval/lecture_retrieval.pyapp/pipeline/lecture_ingestion_pipeline.pyapp/pipeline/shared/citation_pipeline.pyapp/pipeline/chat/course_chat_pipeline.pyapp/pipeline/chat/interaction_suggestion_pipeline.pyapp/pipeline/chat/lecture_chat_pipeline.pyapp/pipeline/chat/exercise_chat_pipeline.pyapp/pipeline/chat/code_feedback_pipeline.py
🔗 Analysis chain
LGTM: Token tracking logic added
The new token tracking logic aligns well with the PR objectives. It correctly captures the token usage and associates it with the specific pipeline type.
However, consider adding error handling to ensure self.llm.tokens is not None before accessing it. Here's a suggested improvement:
if hasattr(self.llm, 'tokens') and self.llm.tokens is not None:
num_tokens = self.llm.tokens
num_tokens.pipeline = PipelineEnum.IRIS_RERANKER_PIPELINE
self.tokens.append(num_tokens)
else:
logger.warning("Token usage information is not available for this LLM.")To verify the usage of self.llm.tokens across the codebase, you can run the following script:
This will help ensure consistent handling of token usage across different parts of the application.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the usage of self.llm.tokens across the codebase
# Test: Search for occurrences of self.llm.tokens
rg --type python -A 5 'self\.llm\.tokens'
Length of output: 74
Script:
#!/bin/bash
# Description: Verify the usage of self.llm.tokens across the codebase by searching Python files
# Test: Search for occurrences of self.llm.tokens in .py files
rg --glob '*.py' -A 5 'self\.llm\.tokens'
Length of output: 5311
| @@ -38,6 +39,7 @@ def __init__(self): | |||
| with open(prompt_file_path, "r") as file: | |||
| self.prompt_str = file.read() | |||
| self.pipeline = self.llm | StrOutputParser() | |||
| self.tokens = [] | |||
There was a problem hiding this comment.
Potential concurrency issues with self.tokens
Initializing self.tokens as a list may lead to concurrency issues if instances of CitationPipeline are used across multiple threads or asynchronous tasks. Appending to a shared list without proper synchronization can cause race conditions. Consider using thread-safe data structures or handling token counts in a way that avoids shared mutable state.
app/retrieval/lecture_retrieval.py
Outdated
| tokens: LLMTokenCount | ||
|
|
There was a problem hiding this comment.
Initialize tokens in the __init__ method
The tokens attribute is declared but not initialized in the constructor. This could lead to an AttributeError if accessed before assignment. It is recommended to initialize tokens in the __init__ method.
Apply this diff to initialize tokens:
def __init__(self, client: WeaviateClient, **kwargs):
super().__init__(implementation_id="lecture_retrieval_pipeline")
+ self.tokens = LLMTokenCount()
request_handler = CapabilityRequestHandler(
requirements=RequirementList(
gpt_version_equivalent=3.5,Committable suggestion was skipped due to low confidence.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Outside diff range and nitpick comments (5)
app/domain/data/token_usage_dto.py (1)
6-12: LGTM: Well-structured DTO for token usage tracking.The
TokenUsageDTOclass is well-designed and aligns perfectly with the PR objectives:
- It includes all necessary attributes for tracking token usage and associated costs.
- The use of specific types (str, int, float) enhances data validation.
- The
pipelineattribute usingPipelineEnumallows for categorization of token usage by pipeline.Consider adding a class-level docstring to provide a brief description of the DTO's purpose and its attributes. This would enhance code readability and maintainability. For example:
class TokenUsageDTO(BaseModel): """ Data Transfer Object for token usage information. Attributes: model_info (str): Information about the model used. num_input_tokens (int): Number of input tokens. cost_per_input_token (float): Cost per input token. num_output_tokens (int): Number of output tokens. cost_per_output_token (float): Cost per output token. pipeline (PipelineEnum): The pipeline associated with this token usage. """ # ... (rest of the class definition)app/domain/pyris_message.py (3)
20-20: LGTM: Well-defined field for cost per input token.The
cost_per_input_tokenfield is correctly implemented with the appropriate type, default value, and alias. It supports the PR objective of tracking usage costs.Consider adding a comment explaining the unit of cost (e.g., USD, EUR) to improve clarity:
cost_per_input_token: float = Field(alias="costPerInputToken", default=0) # Cost in USD per input token
22-22: LGTM: Well-defined field for cost per output token.The
cost_per_output_tokenfield is correctly implemented with the appropriate type, default value, and alias. It complements thecost_per_input_tokenfield and supports comprehensive cost tracking.For consistency with the
cost_per_input_tokenfield, consider adding a comment explaining the unit of cost:cost_per_output_token: float = Field(alias="costPerOutputToken", default=0) # Cost in USD per output token
23-23: LGTM: Useful field for storing model information.The
model_infofield is well-defined and provides valuable context about the model used in token generation. This aligns with the PR objective of enhanced tracking for iris requests.To improve clarity, consider adding a comment explaining the expected format or content of the
model_infofield:model_info: str = Field(alias="modelInfo", default="") # Format: "model_name:version" (e.g., "gpt-3.5-turbo:1")app/llm/external/LLMTokenCount.py (1)
29-33: Consider using a multi-line f-string for improved readability.The
__str__method provides a comprehensive string representation of the object. However, the readability could be improved by using a multi-line f-string.Here's a suggested improvement:
def __str__(self): return f"""\ {self.model_info}: Input: {self.num_input_tokens} tokens, cost {self.cost_per_input_token} Output: {self.num_output_tokens} tokens, cost {self.cost_per_output_token} Pipeline: {self.pipeline}"""This format improves readability and makes it easier to add or modify fields in the future.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (5)
- app/domain/data/token_usage_dto.py (1 hunks)
- app/domain/pyris_message.py (1 hunks)
- app/llm/external/LLMTokenCount.py (1 hunks)
- app/llm/langchain/iris_langchain_chat_model.py (3 hunks)
- app/llm/request_handler/capability_request_handler.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- app/llm/langchain/iris_langchain_chat_model.py
🧰 Additional context used
🔇 Additional comments (9)
app/domain/data/token_usage_dto.py (2)
1-3: LGTM: Appropriate imports for the DTO.The imports are well-chosen:
BaseModelfrom pydantic provides robust data validation for the DTO.PipelineEnumimport aligns with the PR objective of tracking token usage across different pipelines.
1-12: Overall: Excellent implementation of TokenUsageDTOThis new file introduces a well-structured DTO for token usage tracking, which aligns perfectly with the PR objectives. The use of pydantic's BaseModel and specific types for attributes ensures robust data validation. The inclusion of the PipelineEnum allows for effective categorization of token usage across different pipelines.
A minor suggestion for adding a class-level docstring was made to enhance code documentation, but this is not critical for functionality.
Great job on implementing this crucial component for tracking token usage in the Pyris project!
app/domain/pyris_message.py (2)
19-19: LGTM: Appropriate field for tracking input tokens.The
num_input_tokensfield is well-defined with the correct type, default value, and alias. It aligns with the PR objective of tracking token usage for iris requests.
21-21: LGTM: Appropriate field for tracking output tokens.The
num_output_tokensfield is well-defined with the correct type, default value, and alias. It complements thenum_input_tokensfield and supports the PR objective of comprehensive token usage tracking.app/llm/external/LLMTokenCount.py (4)
1-1: LGTM: Import statement is correct and follows best practices.The import statement is well-structured, importing only the specific enum needed from the correct module path.
4-11: LGTM: Well-structured class definition with appropriate attributes.The
LLMTokenCountclass is well-defined with clear and descriptive attribute names. The use of type hints for all attributes is commendable, enhancing code readability and type safety. The inclusion of thePipelineEnumfor thepipelineattribute ensures consistency in pipeline identification.
13-27: LGTM: Well-implemented__init__method.The
__init__method is correctly implemented with:
- Parameters matching class attributes
- Consistent use of type hints
- Proper initialization of all attributes
This implementation ensures that all necessary data is provided when creating an instance of
LLMTokenCount.
1-33: Overall, excellent implementation of theLLMTokenCountclass.This new class aligns well with the PR objectives of tracking token usage for iris requests. The code is clean, well-organized, and follows Python best practices. The consistent use of type hints throughout the file is commendable.
A few key strengths:
- Clear and descriptive attribute names
- Proper use of type hints
- Well-structured
__init__method- Informative
__str__methodThe only suggestion for improvement is to consider using a multi-line f-string in the
__str__method for better readability.Great job on implementing this crucial component for token usage tracking!
app/llm/request_handler/capability_request_handler.py (1)
47-50:⚠️ Potential issueEnsure 'cost_per_input_token' and 'cost_per_output_token' are defined in 'PyrisMessage'
Assigning
cost_per_input_tokenandcost_per_output_tokendirectly to themessageobject assumes that these attributes exist in thePyrisMessageclass. If not, this could lead to anAttributeErrorat runtime. It's important to verify that these attributes are properly defined withinPyrisMessage.Run the following script to verify that these attributes are defined:
If these attributes are not defined, consider adding them to the
PyrisMessageclass:+ class PyrisMessage: + # Existing attributes... + cost_per_input_token: float = 0.0 + cost_per_output_token: float = 0.0
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (5)
app/llm/external/ollama.py (1)
Line range hint
60-73: LGTM! Consider enhancing the function documentation.The changes to
convert_to_iris_messagealign well with the PR objectives of tracking token usage. The new parameters allow for more detailed message conversion, including token counts and model information.Consider updating the function's docstring to reflect the new parameters and their purpose. For example:
def convert_to_iris_message( message: Message, num_input_tokens: int, num_output_tokens: int, model: str ) -> PyrisMessage: """ Convert a Message to a PyrisMessage Args: message (Message): The message to convert num_input_tokens (int): The number of input tokens used num_output_tokens (int): The number of output tokens generated model (str): The model used for generation Returns: PyrisMessage: The converted message with token usage information """app/pipeline/competency_extraction_pipeline.py (1)
111-111: Consider passing only the current execution's token usage.The addition of
tokens=self.tokensto the callback is good for providing token usage data. However,self.tokensaccumulates data across multiple pipeline executions. For status updates, it might be more appropriate to pass only the token usage from the current execution.Consider modifying the line as follows:
self.callback.done(final_result=generated_competencies, tokens=[self.tokens[-1]])This ensures that only the most recent token usage is sent with each status update.
app/pipeline/chat/code_feedback_pipeline.py (1)
147-149: LGTM: Token usage tracking implemented.The implementation of token usage tracking is correct and aligns with the PR objectives. The token count is retrieved from the LLM, the pipeline type is set, and the information is stored in the
self.tokensattribute.Consider adding a comment explaining the purpose of these lines for better code readability:
# Track token usage for this pipeline token_usage = self.llm.tokens token_usage.pipeline = PipelineEnum.IRIS_CODE_FEEDBACK self.tokens = token_usageapp/pipeline/lecture_ingestion_pipeline.py (2)
278-285: Token usage tracking implemented, but TODO remainsThe implementation of token usage tracking is good. However, there's a TODO comment indicating that sending this data to Artemis is still pending.
- The
cleanfunction application improves the quality of the merged content.- Token usage is correctly captured and categorized for the
IRIS_LECTURE_INGESTIONpipeline.- The token usage is properly appended to
self.tokens.Would you like assistance in implementing the logic to send this data to Artemis, or should we create a GitHub issue to track this TODO?
Line range hint
1-339: Summary: Token usage tracking successfully implementedThe changes in this file successfully implement token usage tracking for the lecture ingestion pipeline, aligning well with the PR objectives. Key points:
- The
PipelineEnumimport andself.tokensinitialization set up the foundation for token tracking.- Token usage is captured and categorized correctly in the
merge_page_content_and_image_interpretationmethod.- The previous issue with variable reference has been resolved.
However, there's still a pending TODO to send this data to Artemis. Consider addressing this in a follow-up task or PR to fully complete the token usage tracking feature.
To fully realize the benefits of this token tracking implementation, ensure that the data sent to Artemis is structured in a way that allows for easy analysis and reporting of token usage across different pipelines. This will provide valuable insights into resource utilization and help in optimizing the system's performance.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (10)
- app/llm/external/PipelineEnum.py (1 hunks)
- app/llm/external/ollama.py (3 hunks)
- app/pipeline/chat/code_feedback_pipeline.py (3 hunks)
- app/pipeline/chat/exercise_chat_pipeline.py (9 hunks)
- app/pipeline/chat/lecture_chat_pipeline.py (3 hunks)
- app/pipeline/competency_extraction_pipeline.py (4 hunks)
- app/pipeline/lecture_ingestion_pipeline.py (3 hunks)
- app/pipeline/pipeline.py (1 hunks)
- app/pipeline/shared/reranker_pipeline.py (3 hunks)
- app/retrieval/lecture_retrieval.py (7 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- app/llm/external/PipelineEnum.py
- app/pipeline/chat/lecture_chat_pipeline.py
- app/pipeline/pipeline.py
- app/pipeline/shared/reranker_pipeline.py
🧰 Additional context used
🔇 Additional comments (22)
app/llm/external/ollama.py (2)
116-121: LGTM! Good job addressing the previous review comment.The changes to the
chatmethod effectively implement token usage tracking and address the concerns raised in the previous review about error handling for missing keys. Using thegetmethod with default values is a robust approach to handle potential missing data.This implementation successfully addresses the previous review comment about adding error handling for missing keys in the 'response' dictionary.
Line range hint
60-121: Overall, excellent implementation of token usage tracking!The changes in this file successfully implement token usage tracking and address previous concerns about error handling. The modifications align well with the PR objectives and enhance the robustness of the code. Great job on improving both functionality and code quality!
app/pipeline/competency_extraction_pipeline.py (4)
17-18: LGTM: New imports for token tracking.The new imports for
LLMTokenCountandPipelineEnumare correctly added and necessary for the token tracking functionality being implemented.
43-43: LGTM: New instance variable for token tracking.The addition of
self.tokensas an empty list is appropriate for accumulating token usage data across multiple pipeline executions.
82-88: LGTM: Token tracking implementation.The implementation of token tracking is well-structured:
- Correctly creates an
LLMTokenCountobject with relevant data.- Uses the appropriate
PipelineEnumto identify this pipeline.- Appends the token usage to
self.tokensfor accumulation.This aligns well with the PR objectives for tracking token usage across different pipelines.
Line range hint
1-111: Overall implementation looks good.The changes made to implement token tracking in the
CompetencyExtractionPipelineare well-structured and align with the PR objectives. The token usage is correctly captured and stored. The only suggestion for improvement is to consider passing only the current execution's token usage to the callback for more accurate status updates.app/pipeline/chat/code_feedback_pipeline.py (3)
16-17: LGTM: New imports for token tracking.The new imports for
LLMTokenCountandPipelineEnumare correctly added and align with the PR objectives of tracking token usage across different pipelines.
45-45: LGTM: New attribute for token tracking.The addition of the
tokensattribute of typeLLMTokenCountis appropriate for storing token usage information in theCodeFeedbackPipelineclass.
Line range hint
1-149: Summary: Token usage tracking successfully implemented.The changes in this file successfully implement token usage tracking for the
CodeFeedbackPipelineclass. The new imports, attribute, and logic in the__call__method all contribute to achieving the PR objective of enhancing the tracking of usage related to iris requests.The implementation allows for:
- Storing token usage information specific to this pipeline.
- Setting the correct pipeline type for accurate tracking.
- Preparing the data to be sent to Artemis as part of a status update.
These changes lay the groundwork for improved monitoring and management of token usage across different pipelines in the Pyris project.
app/pipeline/lecture_ingestion_pipeline.py (3)
21-21: LGTM: New import for PipelineEnumThe addition of
PipelineEnumimport is appropriate and aligns with the PR objectives of tracking token usage across different pipelines.
116-116: LGTM: Initialization of self.tokensThe initialization of
self.tokensas an empty list is appropriate for storing token usage information during the pipeline's execution.
284-284: Fixed: Correct usage of self.tokensThe previous issue of using
tokensinstead ofself.tokenshas been addressed. The code now correctly usesself.tokens.append(token_usage), resolving the variable reference problem.app/pipeline/chat/exercise_chat_pipeline.py (6)
37-37: LGTM: Import statement for PipelineEnum addedThe import statement for
PipelineEnumis correctly placed and aligns with the PR objectives of enhancing token usage tracking across different pipelines.
82-82: LGTM: self.tokens initialized for token trackingThe addition of
self.tokens = []in the__init__method is appropriate for accumulating token usage data throughout the pipeline execution, which aligns with the PR objectives.
103-105: LGTM: Token handling improved in call methodThe changes in the
__call__method effectively incorporate token tracking:
- Tokens are now included in the
callback.done()method.- Suggestion pipeline tokens are handled with a null check before being passed to the callback.
These modifications align with the PR objectives and address a previous review comment about null checking
self.suggestion_pipeline.tokens.Also applies to: 119-127
215-216: LGTM: Token collection improved in _run_exercise_chat_pipeline methodThe changes in the
_run_exercise_chat_pipelinemethod enhance token tracking:
- Tokens from
code_feedback_pipelineandretrieverare now appended toself.tokenswith null checks.- The new
_collect_llm_tokens()method is called at appropriate points in the pipeline execution.These modifications align with the PR objectives and address a previous review comment about refactoring duplicated code for collecting LLM tokens.
Also applies to: 237-238, 271-271, 286-286
389-392: LGTM: _collect_llm_tokens method added for centralized token collectionThe new
_collect_llm_tokensmethod effectively centralizes the logic for collecting LLM tokens:
- It includes a null check before appending tokens, which is good practice.
- The method sets the
pipelineattribute of the tokens toPipelineEnum.IRIS_CHAT_EXERCISE_MESSAGE, aligning with the PR objectives.This addition addresses a previous review comment about refactoring duplicated code for collecting LLM tokens, improving code maintainability.
Line range hint
37-392: Overall LGTM: Token tracking successfully implementedThe changes in this file effectively implement token tracking across the
ExerciseChatPipeline:
- Token collection is consistently implemented for various pipeline components (LLM, code feedback, retriever, suggestion pipeline).
- Proper null checks and error handling have been added where necessary.
- Code structure has been improved by refactoring duplicated logic into the
_collect_llm_tokensmethod.These modifications align well with the PR objectives of enhancing the tracking of usage related to iris requests. The implementation will provide better insights into resource utilization across different pipelines.
app/retrieval/lecture_retrieval.py (4)
284-286: Ensure consistent thread-safe token usage trackingThe token usage tracking implementation in this method is consistent with
rewrite_student_query, which is good for maintainability. However, it has the same potential thread safety issue.Please apply the same solution as suggested for the
rewrite_student_querymethod to ensure thread safety.
322-324: Ensure consistent thread-safe token usage tracking and fix typoThe token usage tracking implementation in this method is consistent with previous methods, which is good for maintainability. However, it has the same potential thread safety issue.
Please apply the same solution as suggested for the
rewrite_student_querymethod to ensure thread safety.There's a typo in the logging message:
- logger.info(f"Response from retirval pipeline: {response}") + logger.info(f"Response from retrieval pipeline: {response}")
364-366: Ensure consistent thread-safe token usage trackingThe token usage tracking implementation in this method is consistent with previous methods, which is good for maintainability. However, it has the same potential thread safety issue.
Please apply the same solution as suggested for the
rewrite_student_querymethod to ensure thread safety.
Line range hint
1-524: Summary of changes and recommendationsThe implementation of token usage tracking aligns well with the PR objectives. The changes are consistent across multiple methods, which is good for maintainability. However, there are a few important points to address:
- Initialize the
tokensattribute in the constructor to prevent potentialAttributeErrors.- Implement a thread-safe solution for token usage tracking, as suggested in the previous comments, to avoid race conditions when methods are invoked concurrently.
- Fix the typo in the logging message in the
rewrite_elaborated_querymethod.Addressing these points will significantly improve the reliability and correctness of the token tracking feature. Overall, the changes are a good step towards enhancing the monitoring and management of token usage across different pipelines.
app/retrieval/lecture_retrieval.py
Outdated
| from ..llm.external.LLMTokenCount import LLMTokenCount | ||
| from ..llm.external.PipelineEnum import PipelineEnum |
There was a problem hiding this comment.
Initialize tokens attribute in the constructor
The new imports and the tokens attribute declaration align with the PR objectives for token usage tracking. However, the tokens attribute is not initialized in the constructor, which could lead to AttributeError if accessed before assignment.
To fix this, initialize the tokens attribute in the __init__ method:
def __init__(self, client: WeaviateClient, **kwargs):
super().__init__(implementation_id="lecture_retrieval_pipeline")
+ self.tokens: List[LLMTokenCount] = []
request_handler = CapabilityRequestHandler(
requirements=RequirementList(
gpt_version_equivalent=3.5,
context_length=16385,
privacy_compliance=True,
)
)
# ... rest of the constructor ...Also applies to: 86-87
alexjoham
changed the title
| self.tokens = LLMTokenCount( | ||
| model_info=iris_message.model_info, | ||
| num_input_tokens=iris_message.num_input_tokens, | ||
| cost_per_input_token=iris_message.cost_per_input_token, | ||
| num_output_tokens=iris_message.num_output_tokens, | ||
| cost_per_output_token=iris_message.cost_per_output_token, | ||
| pipeline=PipelineEnum.NOT_SET, | ||
| ) |
There was a problem hiding this comment.
Wouldn't this overwrite existing any token counts when this wrapper is used a second time?
There was a problem hiding this comment.
The token information is saved in every pipeline as soon as the LLM call is done, so the token counts can be overwritten without any loss
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Outside diff range and nitpick comments (12)
app/common/PipelineEnum.py (1)
5-16: LGTM: Comprehensive and well-structured enum members.The enum members are consistently named and cover a wide range of pipeline types, which aligns well with the PR objectives of tracking usage for various IRIS requests. The inclusion of a
NOT_SETmember is a good practice for handling undefined cases.Consider adding a brief class-level docstring to explain the purpose of this enum and how it relates to the IRIS system. This would enhance the code's self-documentation. For example:
class PipelineEnum(str, Enum): """ Enumeration of pipeline types used in the IRIS system. These constants represent various stages or types of pipelines for tracking and managing different IRIS requests. """ # ... (enum members)app/common/token_usage_dto.py (2)
6-12: LGTM: Well-structured DTO with appropriate field definitions.The
TokenUsageDTOclass is well-defined using Pydantic'sBaseModel. The fields are correctly typed and aliased for JSON compatibility, with sensible default values. The use ofPipelineEnumfor thepipelinefield aligns well with the PR objectives.Consider adding a class-level docstring to describe the purpose and usage of this DTO. For example:
class TokenUsageDTO(BaseModel): """ Data Transfer Object for token usage information. This class encapsulates information about token usage, including input and output token counts, associated costs, and the pipeline in which the tokens were used. """ # ... existing field definitions ...
14-18: LGTM: Comprehensive string representation.The
__str__method provides a clear and informative string representation of theTokenUsageDTOobject, including all relevant field values. The use of an f-string and multi-line formatting enhances readability.Consider adding
.valuewhen referencingself.pipelinein the string representation to ensure you're displaying the enum's value rather than potentially its string representation:f"... pipeline: {self.pipeline.value} "This change would ensure consistency in how the pipeline is represented as a string.
app/pipeline/pipeline.py (1)
36-38: LGTM with a suggestion: Initializeself.tokensin the constructor.The
_append_tokensmethod is well-implemented and aligns with the PR objectives. It correctly sets thepipelineattribute of thetokensobject before appending it to the list.However, to prevent potential
AttributeErrors, consider initializingself.tokensin the constructor:def __init__(self, implementation_id=None, **kwargs): self.implementation_id = implementation_id + self.tokens = []This ensures that
self.tokensis always initialized as an empty list when a newPipelineobject is created.app/llm/request_handler/request_handler_interface.py (2)
Line range hint
33-35: Consider using a more specific type for chat messages.The
chatmethod currently useslist[any]for themessagesparameter. To improve type safety and code clarity, consider using a more specific type that represents the structure of chat messages.You could define a
ChatMessagetype or usePyrisMessageif it's suitable. For example:from typing import List def chat(self, messages: List[PyrisMessage], arguments: CompletionArguments) -> PyrisMessage: """Create a completion from the chat messages""" raise NotImplementedError
Line range hint
24-29: Add return type hint for thecompletemethod.The
completemethod is missing a return type hint. Adding it would improve code readability and type checking.Consider updating the method signature as follows:
def complete( self, prompt: str, arguments: CompletionArguments, image: Optional[ImageMessageContentDTO] = None, ) -> str: """Create a completion from the prompt""" raise NotImplementedErrorapp/pipeline/chat/code_feedback_pipeline.py (1)
147-149: LGTM: Token usage tracking implemented correctly.The implementation of token usage tracking aligns well with the PR objectives. The pipeline type is correctly set, and the token usage is retrieved and stored appropriately.
Consider adding error handling to ensure
self.llm.tokensexists before accessing it:if hasattr(self.llm, 'tokens'): token_usage = self.llm.tokens token_usage.pipeline = PipelineEnum.IRIS_CODE_FEEDBACK self.tokens = token_usage else: logger.warning("Token usage information not available")This addition would make the code more robust against potential changes in the
IrisLangchainChatModelimplementation.app/pipeline/lecture_ingestion_pipeline.py (4)
254-256: Token tracking for image interpretationThe
_append_tokensmethod is called after image interpretation. This is good for tracking token usage for this specific operation.However, consider adding error handling around this call to ensure it doesn't break the main flow if there's an issue with token tracking.
response = self.llm_vision.chat( [iris_message], CompletionArguments(temperature=0, max_tokens=512) ) -self._append_tokens( - response.token_usage, PipelineEnum.IRIS_LECTURE_INGESTION -) +try: + self._append_tokens( + response.token_usage, PipelineEnum.IRIS_LECTURE_INGESTION + ) +except Exception as e: + logger.warning(f"Failed to append tokens: {e}")
285-289: Token tracking for content mergingToken tracking is added for the content merging operation. The use of
cleanfunction is good for maintaining consistent output.However, similar to the previous suggestion, consider adding error handling for the token tracking call.
clean_output = clean( (prompt | self.pipeline).invoke({}), bullets=True, extra_whitespace=True ) -self._append_tokens(self.llm.tokens, PipelineEnum.IRIS_LECTURE_INGESTION) +try: + self._append_tokens(self.llm.tokens, PipelineEnum.IRIS_LECTURE_INGESTION) +except Exception as e: + logger.warning(f"Failed to append tokens: {e}") return clean_output
306-306: Token tracking for language detectionToken usage is tracked for the language detection operation. This is consistent with the token tracking implemented in other methods.
As before, consider adding error handling for robustness.
-self._append_tokens(response.token_usage, PipelineEnum.IRIS_LECTURE_INGESTION) +try: + self._append_tokens(response.token_usage, PipelineEnum.IRIS_LECTURE_INGESTION) +except Exception as e: + logger.warning(f"Failed to append tokens: {e}")
Line range hint
1-346: Overall assessment: Token tracking well-implemented with minor improvement suggestionsThe implementation of token tracking in the
LectureIngestionPipelineclass is thorough and consistent. It aligns well with the PR objectives of enhancing the tracking of usage related to iris requests. The token usage is captured for various operations including image interpretation, content merging, and language detection.Key points:
- Token tracking is initialized properly in the constructor.
- Token usage is included in both success and error callbacks.
- The
_append_tokensmethod is used consistently across different operations.Suggestions for improvement:
- Consider adding error handling around the token tracking calls to ensure they don't disrupt the main flow of the program if there's an issue.
- The
_append_tokensmethod is not visible in this file. Ensure it's properly implemented and documented in the appropriate location.Overall, the changes are well-structured and effectively implement the new token tracking feature.
app/retrieval/lecture_retrieval.py (1)
244-246: Consider appendingtoken_usageinstead ofself.llm.tokensAfter assigning
token_usage = self.llm.tokensand updatingtoken_usage.pipeline, you appendself.llm.tokenstoself.tokens. Sincetoken_usageandself.llm.tokensrefer to the same object, appendingtoken_usageimproves code clarity.Apply this diff:
- self.tokens.append(self.llm.tokens) + self.tokens.append(token_usage)Also applies to: 284-286, 322-324, 364-366
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (28)
- app/common/PipelineEnum.py (1 hunks)
- app/common/message_converters.py (1 hunks)
- app/common/pyris_message.py (2 hunks)
- app/common/token_usage_dto.py (1 hunks)
- app/domain/init.py (0 hunks)
- app/domain/chat/chat_pipeline_execution_base_data_dto.py (1 hunks)
- app/domain/chat/chat_pipeline_execution_dto.py (1 hunks)
- app/domain/chat/interaction_suggestion_dto.py (1 hunks)
- app/domain/status/status_update_dto.py (1 hunks)
- app/llm/external/model.py (1 hunks)
- app/llm/external/ollama.py (3 hunks)
- app/llm/external/openai_chat.py (3 hunks)
- app/llm/langchain/iris_langchain_chat_model.py (3 hunks)
- app/llm/request_handler/basic_request_handler.py (1 hunks)
- app/llm/request_handler/capability_request_handler.py (2 hunks)
- app/llm/request_handler/request_handler_interface.py (1 hunks)
- app/pipeline/chat/code_feedback_pipeline.py (3 hunks)
- app/pipeline/chat/course_chat_pipeline.py (5 hunks)
- app/pipeline/chat/exercise_chat_pipeline.py (11 hunks)
- app/pipeline/chat/interaction_suggestion_pipeline.py (4 hunks)
- app/pipeline/chat/lecture_chat_pipeline.py (3 hunks)
- app/pipeline/competency_extraction_pipeline.py (4 hunks)
- app/pipeline/lecture_ingestion_pipeline.py (8 hunks)
- app/pipeline/pipeline.py (2 hunks)
- app/pipeline/shared/citation_pipeline.py (3 hunks)
- app/pipeline/shared/reranker_pipeline.py (3 hunks)
- app/retrieval/lecture_retrieval.py (7 hunks)
- app/web/status/status_update.py (4 hunks)
💤 Files with no reviewable changes (1)
- app/domain/init.py
✅ Files skipped from review due to trivial changes (2)
- app/common/message_converters.py
- app/domain/chat/chat_pipeline_execution_dto.py
🚧 Files skipped from review as they are similar to previous changes (9)
- app/domain/status/status_update_dto.py
- app/llm/request_handler/capability_request_handler.py
- app/pipeline/chat/course_chat_pipeline.py
- app/pipeline/chat/interaction_suggestion_pipeline.py
- app/pipeline/chat/lecture_chat_pipeline.py
- app/pipeline/competency_extraction_pipeline.py
- app/pipeline/shared/citation_pipeline.py
- app/pipeline/shared/reranker_pipeline.py
- app/web/status/status_update.py
🧰 Additional context used
🔇 Additional comments (49)
app/domain/chat/interaction_suggestion_dto.py (2)
Line range hint
9-12: LGTM! Consider verifying PyrisMessage's new capabilities.The
InteractionSuggestionPipelineExecutionDTOclass remains unchanged and correctly uses the importedPyrisMessage. This is good as it maintains backwards compatibility.Given the PR's focus on token usage tracking, it might be beneficial to verify if
PyrisMessagenow includes new token-related attributes or methods. You can use the following script to inspect thePyrisMessageclass:#!/bin/bash # Description: Inspect the PyrisMessage class for new token-related attributes or methods # Test: Search for the PyrisMessage class definition and its contents. Expect: New token-related attributes or methods. ast-grep --lang python --pattern $'class PyrisMessage($_): $$$ 'If
PyrisMessagehas new capabilities, consider whetherInteractionSuggestionPipelineExecutionDTOshould be updated to utilize these new features.
5-5: LGTM! Verify other imports for consistency.The import statement for
PyrisMessagehas been correctly updated to reflect its new location in theapp.common.pyris_messagemodule. This change aligns with good practices for code organization.To ensure consistency across the codebase, please verify that all other imports of
PyrisMessagehave been updated similarly. You can use the following script to check for any remaining old imports:✅ Verification successful
Verification Successful!
All old import statements for
PyrisMessagehave been successfully updated. No remaining old imports found.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for any remaining old imports of PyrisMessage # Test: Search for old import statements. Expect: No results. rg --type python 'from app\.domain import PyrisMessage'Length of output: 783
app/domain/chat/chat_pipeline_execution_base_data_dto.py (2)
Line range hint
1-20: LGTM: Rest of the file is unchanged and consistentThe
ChatPipelineExecutionBaseDataDTOclass and its structure remain unchanged. The usage ofPyrisMessagein thechat_historyfield is still valid after the import change. No further modifications are necessary in this file.
6-6: Verify the new import path for PyrisMessageThe import statement for
PyrisMessagehas been updated fromapp.domain.pyris_messagetoapp.common.pyris_message. This change aligns with the PR objectives of enhancing token usage tracking and suggests a reorganization of the project structure.To ensure this change is consistent across the codebase and doesn't introduce any issues, please run the following verification script:
This script will help ensure that:
- There are no remaining imports from the old path.
- The new import is used consistently across the codebase.
- The
PyrisMessageclass is actually defined in the new location.- There are no potential circular imports introduced by this change.
✅ Verification successful
Import path for
PyrisMessagehas been correctly updated and verified.All checks confirm that:
- There are no remaining imports from the old path.
- The new import is used consistently across the codebase.
PyrisMessageis properly defined in the new location.- No circular imports have been introduced.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the consistency of PyrisMessage import across the codebase # Test 1: Check for any remaining imports from the old path echo "Checking for any remaining old imports:" rg --type python "from app.domain.pyris_message import PyrisMessage" # Test 2: Verify the new import is used consistently echo "Verifying the new import usage:" rg --type python "from app.common.pyris_message import PyrisMessage" # Test 3: Check if PyrisMessage is defined in the new location echo "Checking if PyrisMessage is defined in the new location:" rg --type python "class PyrisMessage" app/common/pyris_message.py # Test 4: Look for any potential circular imports echo "Checking for potential circular imports:" rg --type python "from app.domain" app/common/pyris_message.pyLength of output: 774
Script:
#!/bin/bash # Description: Re-verify the consistency of PyrisMessage import across the codebase without using --type flag # Test 1: Check for any remaining imports from the old path echo "Checking for any remaining old imports:" fd --extension py --exec rg "from app.domain.pyris_message import PyrisMessage" # Test 2: Verify the new import is used consistently echo "Verifying the new import usage:" fd --extension py --exec rg "from app.common.pyris_message import PyrisMessage" # Test 3: Check if PyrisMessage is defined in the new location echo "Checking if PyrisMessage is defined in the new location:" rg "class PyrisMessage" app/common/pyris_message.py # Test 4: Look for any potential circular imports echo "Checking for potential circular imports:" rg "from app.domain" app/common/pyris_message.pyLength of output: 1175
Script:
#!/bin/bash # Description: Check if app/domain/data/message_content_dto.py imports from app.common.pyris_message to detect circular imports echo "Checking if app/domain/data/message_content_dto.py imports from app.common.pyris_message:" rg "from app\.common\.pyris_message import" app/domain/data/message_content_dto.pyLength of output: 273
app/common/pyris_message.py (3)
8-8: LGTM: New import for TokenUsageDTOThe import statement for
TokenUsageDTOis correctly placed and necessary for the newtoken_usagefield. This change aligns with the PR objectives of enhancing token usage tracking.
Line range hint
1-28: Overall assessment: Changes align with PR objectivesThe modifications to
PyrisMessageclass successfully introduce token usage tracking, which is the main goal of this PR. The implementation is clean and follows good practices. The suggested minor improvements (field ordering) and the verification script will help ensure smooth integration of these changes across the codebase.
20-21: LGTM: New token_usage field addedThe addition of the
token_usagefield of typeTokenUsageDTOaligns well with the PR objectives for tracking token usage in iris requests. The use of a default factory is a good practice for mutable default values.Consider moving the
token_usagefield after the existing fields to maintain consistency with the current field order. This would help preserve the existing JSON structure:class PyrisMessage(BaseModel): model_config = ConfigDict(populate_by_name=True) sent_at: datetime | None = Field(alias="sentAt", default=None) sender: IrisMessageRole contents: List[MessageContentDTO] = [] token_usage: TokenUsageDTO = Field(default_factory=TokenUsageDTO)To ensure that all usages of
PyrisMessageare updated to handle the newtoken_usagefield, we can run the following script:This script will help identify areas where
PyrisMessageis used, allowing us to verify if any updates are needed to accommodate the newtoken_usagefield.✅ Verification successful
Verification Successful:
token_usagefield is properly integrated and utilized consistently across the codebase. No issues found.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for PyrisMessage instantiations and usages # Search for PyrisMessage instantiations echo "PyrisMessage instantiations:" rg --type python -n "PyrisMessage\s*\(" -A 5 # Search for PyrisMessage attribute accesses echo "\nPyrisMessage attribute accesses:" rg --type python -n "\.token_usage" -C 2Length of output: 297
Script:
#!/bin/bash # Description: Check for PyrisMessage instantiations and usages in Python files # Search for PyrisMessage instantiations echo "PyrisMessage instantiations:" rg "PyrisMessage\s*\(" --type py -n -A 5 # Search for PyrisMessage attribute accesses echo "\nPyrisMessage attribute accesses:" rg "\.token_usage" --type py -n -C 2Length of output: 10279
app/common/PipelineEnum.py (3)
1-1: LGTM: Import statement is correct and concise.The import of
Enumfrom theenummodule is appropriate for defining an Enum class.
4-4: LGTM: Well-designed class inheritance for string-based enum.The
PipelineEnumclass inherits from bothstrandEnum, which is a good practice for creating string-based enumerations. This approach allows enum members to be used as strings while retaining the benefits of an enumeration, such as type safety and iteration capabilities.
1-16: Overall assessment: Well-implemented PipelineEnum class.This new
PipelineEnumclass is a solid addition that aligns perfectly with the PR objectives of enhancing the tracking of usage related to IRIS requests. The implementation follows best practices for string-based enums in Python, and the comprehensive list of pipeline types will facilitate better management of token usage across different pipelines.The code is clean, concise, and ready for integration into the Pyris project. It provides a strong foundation for the token usage tracking feature described in the PR summary.
app/common/token_usage_dto.py (2)
1-3: LGTM: Imports are appropriate and well-organized.The necessary imports from Pydantic and the custom
PipelineEnumare correctly included. The empty line between different import sources improves readability.
1-18: Great implementation of TokenUsageDTO for enhanced token tracking.This new
TokenUsageDTOclass aligns perfectly with the PR objectives of enhancing token usage tracking for iris requests. The use of Pydantic'sBaseModelensures robust data validation and easy serialization, which will be beneficial when sending token usage data to Artemis as mentioned in the PR summary.The implementation covers all necessary fields for comprehensive token usage tracking, including model information, input/output token counts, associated costs, and the pipeline enum. This structure will provide valuable insights into resource utilization across different pipelines.
Overall, this is a well-structured and thoughtful addition to the project that should significantly improve the monitoring and management of token usage.
app/llm/request_handler/basic_request_handler.py (2)
3-3: LGTM! Verify consistent import usage across the codebase.The import change for
PyrisMessagelooks good and aligns with the restructuring mentioned in the summary. This change doesn't affect the functionality of theBasicRequestHandlerclass.To ensure consistency, let's verify that this import change has been applied across the entire codebase:
✅ Verification successful
Import consistency verified across the codebase. All old imports of
PyrisMessagehave been removed, and new imports are correctly in place.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for any remaining old imports of PyrisMessage and new imports echo "Checking for old imports:" rg --type python "from app.domain import PyrisMessage" echo "Checking for new imports:" rg --type python "from app.common.pyris_message import PyrisMessage"Length of output: 311
Script:
#!/bin/bash # Description: Check for any remaining old imports of PyrisMessage and verify new imports using fd and grep echo "Checking for old imports:" fd --extension py --type f | xargs grep "from app.domain import PyrisMessage" echo "Checking for new imports:" fd --extension py --type f | xargs grep "from app.common.pyris_message import PyrisMessage"Length of output: 1103
Line range hint
10-38: Verify if BasicRequestHandler needs updates for token tracking.The PR objectives mention adding a new
tokensargument to all Pipelines, which is not reflected in this class. While this class might not directly handle token usage tracking, it's worth verifying if any updates are needed here to support the new feature.Let's check if similar request handler classes have been updated with token tracking:
app/pipeline/pipeline.py (3)
2-2: LGTM: Import statement is correct and necessary.The addition of
from typing import Listis appropriate for type hinting the newtokensattribute. This follows Python best practices for type hinting.
4-5: LGTM: New imports and class attribute are correctly implemented.The new imports for
TokenUsageDTOandPipelineEnumare necessary for the token tracking feature. Thetokensattribute is correctly type hinted asList[TokenUsageDTO], which aligns with the PR objectives of tracking token usage across different pipelines.Also applies to: 12-12
Line range hint
1-38: Overall assessment: Changes are well-implemented and align with PR objectives.The modifications to the
Pipelineclass effectively introduce token usage tracking, which aligns with the PR objectives. The new imports, class attribute, and_append_tokensmethod are correctly implemented. The only suggestion for improvement is to initializeself.tokensin the constructor for added robustness.These changes provide a solid foundation for tracking token usage across different pipelines in the Pyris project.
app/llm/langchain/iris_langchain_chat_model.py (4)
1-1: LGTM: New imports added correctly.The new imports for
logging,PipelineEnum, andTokenUsageDTOare correctly added and necessary for the new functionality introduced in this update.Also applies to: 11-11, 16-16
25-26: LGTM: New class attributes added correctly.The new
tokensattribute of typeTokenUsageDTOand the logger instantiation are correctly implemented. These additions align with the PR objectives for token usage tracking and logging capabilities.
Line range hint
1-64: Summary: Good implementation with a minor consideration.The changes in this file successfully implement token usage tracking as per the PR objectives. The new imports, class attributes, and modifications to the
_generatemethod are well-integrated and correctly implemented.However, there's a potential issue with overwriting token counts when the
_generatemethod is called multiple times. Consider addressing this by either:
- Accumulating token usage across multiple calls.
- Resetting token usage at an appropriate point in the workflow.
- Documenting the current behavior if it's intentional.
Once this consideration is addressed, the implementation will be robust and fully aligned with the PR objectives.
50-57: LGTM with a consideration: Token usage tracking implementation.The implementation of token usage tracking looks good and aligns with the PR objectives. The
tokensattribute is correctly populated with relevant information fromiris_message.token_usage.However, as pointed out in a previous review comment, there's a potential issue:
If the
_generatemethod is called multiple times on the same instance, the existing token counts would be overwritten. Consider accumulating the token usage instead of overwriting it, or resetting it at an appropriate point in the workflow.app/llm/external/ollama.py (4)
9-10: LGTM: New imports align with PR objectivesThe addition of
PyrisMessageandTokenUsageDTOimports is consistent with the PR's goal of enhancing token usage tracking. These imports are necessary for the subsequent changes in the file.
61-63: LGTM: Function signature update enhances token trackingThe updated signature of
convert_to_iris_messagenow includesnum_input_tokens,num_output_tokens, andmodelparameters. This change aligns well with the PR's objective of enhancing token usage tracking for iris requests.
68-77: LGTM: Token usage tracking implemented correctlyThe changes to the
convert_to_iris_messagefunction body effectively implement token usage tracking. The creation of aTokenUsageDTOinstance and its inclusion in thePyrisMessagealign perfectly with the PR's objectives.
120-125: LGTM: Safe key access implemented and token tracking enhancedThe changes to the
chatmethod effectively address the previous review comment about potentialKeyErrors. The use ofresponse.get()with default values is a safe approach to handle potentially missing keys. Additionally, these changes align with the PR objectives by passing token usage information toconvert_to_iris_message.app/pipeline/chat/code_feedback_pipeline.py (4)
11-17: LGTM: Import changes align with new token tracking functionality.The updated import for
PyrisMessageand the new imports forTokenUsageDTOandPipelineEnumare consistent with the PR objectives of introducing token usage tracking.
45-45: LGTM: Newtokensattribute added for token usage tracking.The addition of the
tokensattribute of typeTokenUsageDTOis consistent with the PR objectives and enables token usage tracking within the pipeline.
47-47: LGTM: Improved method signature with explicit optional parameter.The updated
__init__method signature now explicitly defines thecallbackparameter as optional, enhancing code clarity and type hinting without changing functionality.
Line range hint
1-151: Overall assessment: Changes successfully implement token usage tracking.The modifications to the
CodeFeedbackPipelineclass effectively introduce token usage tracking as per the PR objectives. The changes are well-integrated, maintaining code consistency and enhancing functionality. The newtokensattribute, updated method signatures, and token usage handling in the__call__method all contribute to achieving the goal of tracking and categorizing token usage in the IRIS system.A minor suggestion for error handling has been provided to improve robustness. Otherwise, the implementation looks solid and ready for integration.
app/pipeline/lecture_ingestion_pipeline.py (4)
13-13: New imports added for token trackingThe new imports for
PyrisMessage,IrisMessageRole, andPipelineEnumare related to the token tracking feature. This is a good practice as it keeps the dependencies clear and localized.Also applies to: 21-21
116-116: Token tracking attribute addedThe
self.tokenslist is initialized in the constructor. This is a good place to initialize it, ensuring it's ready for use throughout the class's lifecycle.
144-144: Token usage included in callbackThe
tokensparameter is now included in both successful and error callbacks. This is consistent and allows for proper tracking of token usage regardless of the outcome.Also applies to: 153-155
177-179: Token usage included in error callbackSimilar to the previous comment, tokens are included in the error callback here. This maintains consistency in error reporting.
app/pipeline/chat/exercise_chat_pipeline.py (14)
27-27: ImportPyrisMessagefor message handlingThe import of
PyrisMessageensures proper handling of message types within the pipeline, which is essential for processing chat interactions.
37-37: IntroducePipelineEnumfor pipeline trackingImporting
PipelineEnumallows for better enumeration and tracking of different pipeline stages, enhancing clarity and maintainability.
82-82: Initializeself.tokensfor token aggregationInitializing
self.tokensas an empty list in the constructor correctly sets up the aggregation of token usage data throughout the pipeline execution.
103-105: Pass aggregated tokens to callbackIncluding
tokens=self.tokensin theself.callback.donemethod ensures that token usage data is properly passed back for tracking and monitoring purposes.
119-127: Handlesuggestion_pipeline.tokenssafelyThe code now checks if
self.suggestion_pipeline.tokensis notNonebefore using it, preventing potential errors by avoiding passingNonevalues. This addresses previous concerns about potentialNonevalues being passed.
140-142: Include tokens in error callback for suggestionsPassing
tokens=self.tokensin the error handling of interaction suggestions provides additional context, aiding in debugging and monitoring token usage even when errors occur.
146-148: Enhance error reporting with token dataIncluding
tokens=self.tokensin the error callback improves error reporting by providing token usage information, which can be valuable for diagnosing issues.
219-220: Collect tokens fromcode_feedback_pipelineProperly checking if
self.code_feedback_pipeline.tokensis notNonebefore appending ensures that only valid token data is collected, preventing potentialNoneTypeerrors.
233-235: Provide tokens in error callback for code feedbackIncluding
tokens=self.tokensin the error handling when retrieving code feedback enhances the ability to debug and trace token usage during failures.
243-244: Aggregate tokens fromretrievercomponentThe code correctly checks for
Nonebefore appendingself.retriever.tokenstoself.tokens, ensuring robust handling of token data from the retrieval process.
251-253: Token-inclusive error handling for lecture retrievalPassing
tokens=self.tokensin the error callback when lecture retrieval fails adds valuable context for debugging and monitoring.
279-281: Refactor LLM token collection with_append_tokensBy using the
_append_tokensmethod to collect tokens fromself.llm, the code eliminates duplication and adheres to the DRY principle, addressing previous refactoring suggestions.
296-298: Consistent use of_append_tokensfor LLM tokensThe consistent application of
_append_tokensfor collecting tokens fromself.llmenhances code maintainability and readability.
307-309: Enhance response creation error handling with tokensIncluding
tokens=self.tokenswhen handling exceptions during response creation provides comprehensive debugging information for token usage.app/retrieval/lecture_retrieval.py (2)
9-11: Imports added align with PR objectivesThe new imports of
TokenUsageDTOandPipelineEnumare appropriate for token usage tracking and pipeline enumeration as described in the PR objectives.
105-105: Initialization oftokensattribute approvedThe
tokensattribute is properly initialized in the constructor.
| @@ -1,7 +1,7 @@ | |||
| from abc import ABCMeta, abstractmethod | |||
| from typing import Optional | |||
|
|
|||
| from ...domain import PyrisMessage | |||
| from ...common.pyris_message import PyrisMessage | |||
There was a problem hiding this comment.
💡 Codebase verification
Inconsistent import paths for PyrisMessage detected.
app/pipeline/lecture_ingestion_pipeline.pyusesfrom ..common.pyris_message import PyrisMessage, IrisMessageRoleapp/retrieval/lecture_retrieval.pyusesfrom ..common.pyris_message import PyrisMessage
Please update these import statements to match the standardized patterns:
- Relative:
from ...common.pyris_message import PyrisMessage - Absolute:
from app.common.pyris_message import PyrisMessage
This will ensure consistency across the codebase and prevent potential import issues.
🔗 Analysis chain
Import path update looks good.
The import statement for PyrisMessage has been updated from a relative import from the domain module to a relative import from the common.pyris_message module. This change is consistent with the PR objectives and the AI-generated summary, which mentioned updates to import paths for PyrisMessage to reflect a new organizational structure.
To ensure this change is consistent across the codebase, let's verify the usage of PyrisMessage:
[/run_scripts]
[/run_scripts]
[/run_scripts]
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the import path for PyrisMessage across the codebase
# Test: Search for PyrisMessage imports
rg --type python "from.*import.*PyrisMessage"
# Test: Search for PyrisMessage usage
rg --type python "PyrisMessage"
Length of output: 141
Script:
#!/bin/bash
# Description: Verify the import path for PyrisMessage across the codebase
# Test: Search for PyrisMessage imports
rg --glob "*.py" "from.*import.*PyrisMessage"
# Test: Search for PyrisMessage usage
rg --glob "*.py" "PyrisMessage"
Length of output: 8395
Script:
#!/bin/bash
# Description: Identify inconsistent import paths for PyrisMessage
# Define acceptable import patterns
acceptable_patterns="from ...common.pyris_message import PyrisMessage|from app.common.pyris_message import PyrisMessage"
# Search for PyrisMessage imports that do NOT match the acceptable patterns
rg --glob "*.py" "from\s+(?!\.\.\.common\.pyris_message|app\.common\.pyris_message)\w+.*import\s+PyrisMessage"
Length of output: 537
Script:
#!/bin/bash
# Description: Identify inconsistent import paths for PyrisMessage without using look-around
# Step 1: Find all PyrisMessage import statements
rg --glob "*.py" "from.*import.*PyrisMessage" > all_pyris_imports.txt
# Step 2: Exclude acceptable import patterns
rg -v -e "from \.\.\.common\.pyris_message import PyrisMessage" -e "from app\.common\.pyris_message import PyrisMessage" all_pyris_imports.txt
Length of output: 383
| @@ -1,7 +1,7 @@ | |||
| from abc import ABCMeta, abstractmethod | |||
| from pydantic import BaseModel | |||
|
|
|||
| from ...domain import PyrisMessage | |||
| from ...common.pyris_message import PyrisMessage | |||
There was a problem hiding this comment.
💡 Codebase verification
Issues Found: Inconsistent Import Paths for PyrisMessage
The import statements for PyrisMessage are not fully consistent across the codebase:
- Relative imports still exist:
from ..common.pyris_message import PyrisMessage
- Absolute imports are used in some files:
from app.common.pyris_message import PyrisMessage
Please standardize all import paths to from ...common.pyris_message import PyrisMessage to ensure consistency and prevent potential issues.
🔗 Analysis chain
LGTM! Verify import consistency across the codebase.
The import statement for PyrisMessage has been updated to reflect the new module structure. This change is in line with the PR objectives.
To ensure consistency across the codebase, please run the following script:
This script will help identify any inconsistencies in the import statements across the project.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the consistency of PyrisMessage import across the codebase.
# Test 1: Check for any remaining imports from the old path
echo "Checking for any remaining old imports:"
rg --type python "from \.\.\.domain import PyrisMessage"
# Test 2: Verify the new import is used consistently
echo "Verifying the new import usage:"
rg --type python "from \.\.\.common\.pyris_message import PyrisMessage"
# Test 3: Check for any other variations of PyrisMessage import
echo "Checking for any other variations of PyrisMessage import:"
rg --type python "import.*PyrisMessage" | grep -v "from ...common.pyris_message import PyrisMessage"
Length of output: 615
Script:
#!/bin/bash
# Description: Verify the consistency of PyrisMessage import across the codebase.
# Test 1: Check for any remaining imports from the old path
echo "Checking for any remaining old imports:"
rg --type py "from \.\.\.domain import PyrisMessage"
# Test 2: Verify the new import is used consistently
echo "Verifying the new import usage:"
rg --type py "from \.\.\.common\.pyris_message import PyrisMessage"
# Test 3: Check for any other variations of PyrisMessage import
echo "Checking for any other variations of PyrisMessage import:"
rg --type py "import.*PyrisMessage" | grep -v "from ...common.pyris_message import PyrisMessage"
Length of output: 2314
| sender=map_str_to_role(message.role), | ||
| contents=[TextMessageContentDTO(textContent=message.content)], | ||
| send_at=datetime.now(), | ||
| sentAt=datetime.now(), |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use UTC timestamps to ensure consistency across time zones
In the construction of PyrisMessage, consider using datetime.utcnow() instead of datetime.now() to ensure that the timestamp is in UTC. This helps prevent issues related to timezone differences and ensures consistency across different environments.
Apply this diff to update the timestamp:
- sentAt=datetime.now(),
+ sentAt=datetime.utcnow(),📝 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.
| sentAt=datetime.now(), | |
| sentAt=datetime.utcnow(), |
| def convert_to_iris_message( | ||
| message: ChatCompletionMessage, usage: Optional[CompletionUsage], model: str | ||
| ) -> PyrisMessage: |
There was a problem hiding this comment.
💡 Codebase verification
Calls to convert_to_iris_message in app/llm/external/ollama.py do not match the updated function signature. Please update these calls to include all required parameters: message, num_input_tokens, num_output_tokens, and model.
app/llm/external/ollama.py:- Line where
convert_to_iris_messageis called with insufficient arguments.
- Line where
🔗 Analysis chain
Ensure all calls to convert_to_iris_message are updated with new parameters
Since the function convert_to_iris_message now includes additional parameters usage: Optional[CompletionUsage] and model: str, please verify that all calls to this function across the codebase have been updated to match the new signature to prevent any potential errors.
You can run the following script to identify any calls that may need to be updated:
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find all usages of `convert_to_iris_message` and check if they match the new signature.
# Expected: All calls should include three parameters: message, usage, and model.
rg --type python $'convert_to_iris_message\s*\(' -A 1
Length of output: 85
Script:
#!/bin/bash
# Description: Find all usages of `convert_to_iris_message` in Python files and check if they match the new signature.
# Expected: All calls should include three parameters: message, usage, and model.
rg --glob "*.py" 'convert_to_iris_message\s*\(' -A 1
Length of output: 723
app/retrieval/lecture_retrieval.py
Outdated
| tokens: [TokenUsageDTO] | ||
|
|
There was a problem hiding this comment.
Correct the type annotation for tokens attribute
The type annotation tokens: [TokenUsageDTO] is incorrect. It should use List[TokenUsageDTO] from the typing module.
Apply this diff to fix the type annotation:
- tokens: [TokenUsageDTO]
+ tokens: List[TokenUsageDTO]📝 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.
| tokens: [TokenUsageDTO] | |
| tokens: List[TokenUsageDTO] | |
| ``` | |
| Note: This suggestion assumes that the `List` type is imported from the `typing` module elsewhere in the file. If it's not already imported, you would need to add the following import at the top of the file: | |
| ```python | |
| from typing import List |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
app/llm/external/ollama.py (1)
61-77: LGTM with a minor suggestion: Function update aligns with PR objectivesThe changes to
convert_to_iris_messagecorrectly implement token usage tracking, which aligns with the PR objectives. The new parameters are used appropriately to create aTokenUsageDTOobject.For consistency, consider updating the
TokenUsageDTOfield names to use snake_case:tokens = TokenUsageDTO( - numInputTokens=num_input_tokens, - numOutputTokens=num_output_tokens, + num_input_tokens=num_input_tokens, + num_output_tokens=num_output_tokens, modelInfo=model, )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
- app/llm/external/ollama.py (3 hunks)
🧰 Additional context used
🔇 Additional comments (3)
app/llm/external/ollama.py (3)
9-10: LGTM: New imports align with PR objectivesThe addition of
PyrisMessageandTokenUsageDTOimports is consistent with the PR's goal of enhancing token usage tracking.
120-125: LGTM: Improved error handling and consistency with other implementationsThe changes to the
chatmethod effectively implement token usage tracking and address previous concerns:
- Using
getwith default values mitigates potentialKeyErrorissues.- The default value for token counts is now -1, consistent with the OpenAI implementation as mentioned in the past review comments.
These updates align well with the PR objectives and improve the overall robustness of the code.
Line range hint
1-135: Overall assessment: Well-implemented token usage trackingThe changes in this file successfully implement token usage tracking for the Ollama integration, aligning well with the PR objectives. The code quality is good, with only minor consistency suggestions. The implementation addresses previous concerns and improves error handling.
Great job on enhancing the Ollama integration with token usage tracking!
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Outside diff range and nitpick comments (1)
app/llm/external/ollama.py (1)
61-77: LGTM: Token usage tracking implemented correctlyThe changes to the
convert_to_iris_messagefunction successfully implement token usage tracking, aligning with the PR objectives. The new parameters are correctly used to create aTokenUsageDTOobject, which is then included in thePyrisMessage.Minor suggestion:
Consider using snake_case for consistency in theTokenUsageDTOinitialization:tokens = TokenUsageDTO( num_input_tokens=num_input_tokens, num_output_tokens=num_output_tokens, model=model, )This would align with Python's naming conventions and the style used in the rest of the code.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
- app/common/token_usage_dto.py (1 hunks)
- app/llm/external/ollama.py (3 hunks)
- app/llm/external/openai_chat.py (4 hunks)
- app/llm/langchain/iris_langchain_chat_model.py (3 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- app/common/token_usage_dto.py
- app/llm/langchain/iris_langchain_chat_model.py
🧰 Additional context used
🔇 Additional comments (9)
app/llm/external/ollama.py (3)
9-10: LGTM: New imports align with PR objectivesThe addition of
PyrisMessageandTokenUsageDTOimports is consistent with the PR's goal of introducing token usage tracking. These imports are necessary for the subsequent changes in the file.
120-125: LGTM: Safe key access and consistent default valuesThe changes to the
chatmethod effectively implement token usage tracking and address previous review comments:
- The use of the
getmethod with default values prevents potentialKeyErrorissues, as suggested in a past review.- The default value of -1 for token counts is now consistent with other models, addressing the concern raised in another past comment.
These modifications align well with the PR objectives and improve the robustness of the code.
Line range hint
1-135: Summary: Successful implementation of token usage trackingThe changes in this file effectively implement token usage tracking for the Ollama model, aligning well with the PR objectives. Key improvements include:
- Addition of necessary imports for
PyrisMessageandTokenUsageDTO.- Update of the
convert_to_iris_messagefunction to include token usage information.- Modification of the
chatmethod to safely retrieve and pass token usage data.These changes enhance the functionality of the Ollama model integration while addressing previous review comments. The code is now more robust and consistent with other parts of the system.
Great job on implementing these features and addressing the feedback from previous reviews!
app/llm/external/openai_chat.py (6)
4-4: LGTM: Import statements updated correctlyThe new import statements for
Optional,CompletionUsage,PyrisMessage, andTokenUsageDTOare consistent with the PR objectives and support the implementation of token usage tracking.Also applies to: 15-15, 21-22
72-74: LGTM: Function signature updated appropriatelyThe
convert_to_iris_messagefunction signature has been correctly updated to includeusage: Optional[CompletionUsage]andmodel: strparameters. This change aligns with the PR objectives and enables token usage tracking.
81-85: LGTM: TokenUsageDTO creation implemented correctlyThe creation of the
TokenUsageDTOobject with model information and token counts is implemented correctly. This change aligns with the PR objectives of enhancing token usage tracking.
87-91: LGTM with a suggestion: Use UTC time for consistencyThe addition of
token_usageto thePyrisMessagecreation is correct and aligns with the PR objectives. However, as previously suggested, consider usingdatetime.utcnow()instead ofdatetime.now()for thesentAtattribute to ensure consistency across different time zones.Apply this diff to update the timestamp:
- sentAt=datetime.now(), + sentAt=datetime.utcnow(),
131-132: LGTM: Correct extraction of usage and model informationThe extraction of
usageandmodelinformation from the response is implemented correctly. These changes support the updatedconvert_to_iris_messagefunction and align with the PR objectives of enhancing token usage tracking.
139-139: LGTM: Correct usage of updated convert_to_iris_message functionThe
convert_to_iris_messagefunction call has been correctly updated to include theusageandmodelparameters. This change is consistent with the updated function signature and ensures that token usage information is properly passed to the PyrisMessage creation.
app/llm/external/openai_chat.py
Outdated
| num_input_tokens = getattr(usage, "prompt_tokens", -1) | ||
| num_output_tokens = getattr(usage, "completion_tokens", -1) |
There was a problem hiding this comment.
Consider using 0 as the default value for token counts
The current implementation uses -1 as the default value for num_input_tokens and num_output_tokens. However, this might be confusing as it's not a valid token count. For consistency with other parts of the codebase (e.g., ollama) and improved clarity, consider using 0 as the default value.
Suggested change:
- num_input_tokens = getattr(usage, "prompt_tokens", -1)
- num_output_tokens = getattr(usage, "completion_tokens", -1)
+ num_input_tokens = getattr(usage, "prompt_tokens", 0)
+ num_output_tokens = getattr(usage, "completion_tokens", 0)This change would make the default values more intuitive and consistent with other parts of the codebase.
📝 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.
| num_input_tokens = getattr(usage, "prompt_tokens", -1) | |
| num_output_tokens = getattr(usage, "completion_tokens", -1) | |
| num_input_tokens = getattr(usage, "prompt_tokens", 0) | |
| num_output_tokens = getattr(usage, "completion_tokens", 0) |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (1)
app/llm/external/ollama.py (1)
61-63: Add type hints for token count parameters.The function parameters
num_input_tokensandnum_output_tokensshould have explicit type hints for better code maintainability and IDE support.def convert_to_iris_message( - message: Message, num_input_tokens: int, num_output_tokens: int, model: str + message: Message, num_input_tokens: Optional[int], num_output_tokens: Optional[int], model: str ) -> PyrisMessage:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
- app/llm/external/ollama.py (3 hunks)
- app/llm/external/openai_chat.py (4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- app/llm/external/openai_chat.py
🧰 Additional context used
🔇 Additional comments (1)
app/llm/external/ollama.py (1)
9-10: LGTM: New imports align with token tracking requirements.The added imports for
PyrisMessageandTokenUsageDTOare appropriately placed and necessary for the token tracking implementation.
| tokens = TokenUsageDTO( | ||
| numInputTokens=num_input_tokens, | ||
| numOutputTokens=num_output_tokens, | ||
| model=model, | ||
| ) |
There was a problem hiding this comment.
Add validation for token count parameters.
The token counts should be validated to ensure they are not negative values before creating the TokenUsageDTO.
+ if num_input_tokens is not None and num_input_tokens < -1:
+ raise ValueError("Input token count cannot be less than -1")
+ if num_output_tokens is not None and num_output_tokens < -1:
+ raise ValueError("Output token count cannot be less than -1")
tokens = TokenUsageDTO(
numInputTokens=num_input_tokens,
numOutputTokens=num_output_tokens,
model=model,
)Committable suggestion was skipped due to low confidence.
| return convert_to_iris_message( | ||
| response.get("message"), | ||
| response.get("prompt_eval_count", 0), | ||
| response.get("eval_count", 0), | ||
| response.get("model", self.model), | ||
| ) |
There was a problem hiding this comment.
Update default token counts to -1 for consistency.
As discussed in previous comments and confirmed by you, the default token count values should be -1 to maintain consistency with the OpenAI implementation.
return convert_to_iris_message(
response.get("message"),
- response.get("prompt_eval_count", 0),
- response.get("eval_count", 0),
+ response.get("prompt_eval_count", -1),
+ response.get("eval_count", -1),
response.get("model", self.model),
)📝 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.
| return convert_to_iris_message( | |
| response.get("message"), | |
| response.get("prompt_eval_count", 0), | |
| response.get("eval_count", 0), | |
| response.get("model", self.model), | |
| ) | |
| return convert_to_iris_message( | |
| response.get("message"), | |
| response.get("prompt_eval_count", -1), | |
| response.get("eval_count", -1), | |
| response.get("model", self.model), | |
| ) |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Outside diff range and nitpick comments (1)
app/pipeline/chat/exercise_chat_pipeline.py (1)
82-82: Consider adding debug logging for token trackingTo aid in debugging token-related issues, consider adding debug logging when tokens are collected or appended.
Add logging statements like:
def _append_tokens(self, tokens, source): + logger.debug(f"Appending {len(tokens)} tokens from {source}") # existing implementation
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
- app/pipeline/chat/exercise_chat_pipeline.py (11 hunks)
- app/retrieval/lecture_retrieval.py (7 hunks)
🧰 Additional context used
🔇 Additional comments (5)
app/pipeline/chat/exercise_chat_pipeline.py (4)
37-37: LGTM: Token tracking initializationThe initialization of token tracking is properly implemented with the required import and class variable setup.
Also applies to: 82-82
103-105: LGTM: Token tracking in callbackThe token tracking is properly integrated into the callback mechanism.
140-142: LGTM: Comprehensive error handling with token trackingToken tracking is consistently implemented across all error handlers, ensuring token usage is properly reported even in failure scenarios.
Also applies to: 146-148, 233-235, 254-256, 310-312
282-284: Show implementation of_append_tokensThe code uses
_append_tokensmethod but its implementation is not visible in the provided code. Please ensure this method is properly implemented to handle token collection and pipeline assignment.Let's verify if the method is implemented:
Also applies to: 299-301
✅ Verification successful
Let me try a different approach to search for the implementation of
_append_tokens.
Implementation of
_append_tokensis properly defined in the base pipeline classThe method
_append_tokensis implemented inapp/pipeline/pipeline.pyas:def _append_tokens(self, tokens: TokenUsageDTO, pipeline: PipelineEnum) -> None: tokens.pipeline = pipeline self.tokens.append(tokens)The implementation is correct and properly handles:
- Setting the pipeline type for the tokens
- Appending the tokens to the pipeline's token collection
The usage in the exercise chat pipeline is consistent with this implementation and other pipeline classes.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for _append_tokens implementation ast-grep --pattern 'def _append_tokens($_, $_) { $$$ }'Length of output: 57
Script:
#!/bin/bash # Search for _append_tokens using ripgrep for better coverage rg "_append_tokens" -A 5Length of output: 5052
app/retrieval/lecture_retrieval.py (1)
9-10: LGTM! Import statements and class attribute initialization look good.The imports are correctly organized, the type annotation uses
Listfrom the typing module, and thetokenslist is properly initialized in the constructor.Also applies to: 86-87, 105-105
| token_usage = self.llm.tokens | ||
| token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE | ||
| self.tokens.append(self.llm.tokens) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Refactor token tracking to address thread safety and code duplication.
The current implementation has two issues:
- Token tracking code is duplicated across methods
- Concurrent access to
self.tokenscan lead to race conditions when methods run in parallel
Consider this comprehensive solution:
- First, extract the token tracking logic into a helper method:
def _track_token_usage(self) -> TokenUsageDTO:
"""Track token usage for the current operation."""
token_usage = self.llm.tokens
token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
return token_usage- Then, modify the methods to return both the response and token usage:
def rewrite_student_query(self, ...):
try:
response = (prompt | self.pipeline).invoke({})
- token_usage = self.llm.tokens
- token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
- self.tokens.append(self.llm.tokens)
+ token_usage = self._track_token_usage()
logger.info(f"Response from exercise chat pipeline: {response}")
- return response
+ return response, token_usage- Finally, update
run_parallel_rewrite_tasksto safely aggregate tokens:
def run_parallel_rewrite_tasks(self, ...):
with concurrent.futures.ThreadPoolExecutor() as executor:
# ... existing code ...
- rewritten_query = rewritten_query_future.result()
- hypothetical_answer_query = hypothetical_answer_query_future.result()
+ rewritten_query, tokens1 = rewritten_query_future.result()
+ hypothetical_answer_query, tokens2 = hypothetical_answer_query_future.result()
+ # Safely aggregate tokens after all tasks complete
+ self.tokens.extend([tokens1, tokens2])This solution:
- Eliminates code duplication
- Ensures thread safety by aggregating tokens after all tasks complete
- Improves maintainability through centralized token tracking logic
Also applies to: 284-286, 322-324, 364-366
Track usage of iris requests
All Pipelines now have a tokens argument that is set by langchain. The tokens have a new type called LLMTokenCount, where an Enum is also added to set the pipeline in which the specified cost was created. Pyris sends now with an status update the tokens to Artemis.
New class LLMTokenCount with the following attributes:
Every LLM call has those attributes set by langchain to track the usage. pipeline is an enum with all the important pipelines that Pyris has:
Langchain sets the enum to not_set and the pipeline edits the variable and sets it to their value.
Summary by CodeRabbit
Summary by CodeRabbit
New Features
TokenUsageDTOfor tracking token usage information.PyrisMessageto include token counts and model info.PipelineEnumfor managing different processing stages.CapabilityRequestHandlerto return messages with cost information for input and output tokens.IrisLangchainChatModelto incorporate token counting functionality.CodeFeedbackPipeline,ExerciseChatPipeline, andLectureChatPipeline.LectureRetrieval,CitationPipeline, andCompetencyExtractionPipeline.StatusUpdateDTOto include token usage information.doneanderrormethods in theStatusCallbackclass to accept token usage parameters.Bug Fixes
Documentation