From 12de4cfb73c86ebec297eccb759ef1f793c8f6fb Mon Sep 17 00:00:00 2001 From: Saurav Panda Date: Fri, 31 May 2024 12:55:46 -0400 Subject: [PATCH] bugfix: work summary wouldnt process the file if available token > 0 --- kaizen/llms/provider.py | 66 ++++++++++++++++++++++++++++++ kaizen/reviewer/code_review.py | 4 +- kaizen/reviewer/work_summarizer.py | 11 +++++ 3 files changed, 80 insertions(+), 1 deletion(-) diff --git a/kaizen/llms/provider.py b/kaizen/llms/provider.py index b61b5787..ee79ee39 100644 --- a/kaizen/llms/provider.py +++ b/kaizen/llms/provider.py @@ -45,6 +45,16 @@ def is_inside_token_limit(self, PROMPT, percentage=0.7): return True def available_tokens(self, message, percentage=0.8): + """ + Calculate the number of tokens available for a single request after accounting for the tokens consumed by the prompt. + + Args: + message (str): The input message for which tokens are being calculated. + percentage (float, optional): The percentage of total tokens to be considered available. Defaults to 0.8. + + Returns: + int: The number of tokens available for a single request. + """ return litellm.get_max_tokens(self.model) * percentage - litellm.token_counter( model=self.model, text=message ) @@ -53,6 +63,43 @@ def get_token_count(self, message): return litellm.token_counter(model=self.model, text=message) def update_usage(self, total_usage, current_usage): + """ + Update the total usage with the current usage values. + + This function updates the `total_usage` dictionary by adding the values from the + `current_usage` dictionary. If `total_usage` is None, it initializes `total_usage` with + the values from `current_usage`. + + Args: + total_usage (dict or None): The dictionary containing the cumulative usage information. + If None, it will be initialized with the values from `current_usage`. + current_usage (dict): A dictionary containing the current usage information with keys + corresponding to those in `total_usage`. + + Returns: + dict: The updated total usage dictionary with the values from `current_usage` added + to the corresponding keys in `total_usage`. + + Example: + total_usage = { + "prompt_tokens": 150, + "completion_tokens": 250 + } + current_usage = { + "prompt_tokens": 100, + "completion_tokens": 200 + } + updated_usage = instance.update_usage(total_usage, current_usage) + print(updated_usage) + # Output: {"prompt_tokens": 250, "completion_tokens": 450} + + # If total_usage is None + total_usage = None + updated_usage = instance.update_usage(total_usage, current_usage) + print(updated_usage) + # Output: {"prompt_tokens": 100, "completion_tokens": 200} + """ + if total_usage is not None: total_usage = { key: total_usage[key] + current_usage[key] for key in total_usage @@ -62,6 +109,25 @@ def update_usage(self, total_usage, current_usage): return total_usage def get_usage_cost(self, total_usage): + """ + Calculate the cost of usage based on the number of tokens used in the prompt and completion. + + Args: + total_usage (dict): A dictionary containing the usage information with the following keys: + - "prompt_tokens" (int): The number of tokens used in the prompt. + - "completion_tokens" (int): The number of tokens used in the completion. + + Returns: + float: The total cost of the usage based on the model's cost per token. + + Example: + total_usage = { + "prompt_tokens": 100, + "completion_tokens": 200 + } + cost = instance.get_usage_cost(total_usage) + print(cost) # Output will depend on the model's cost per token. + """ return litellm.cost_per_token( self.model, total_usage["prompt_tokens"], total_usage["completion_tokens"] ) diff --git a/kaizen/reviewer/code_review.py b/kaizen/reviewer/code_review.py index 7e673d59..f0928952 100644 --- a/kaizen/reviewer/code_review.py +++ b/kaizen/reviewer/code_review.py @@ -166,7 +166,9 @@ def create_pr_review_text(self, topics): markdown_output += ct + "\n" if high_ranked_issues > 0: - status_msg = "❗ **Attention Required:** This PR has potential issues. 🚨\n\n" + status_msg = ( + "❗ **Attention Required:** This PR has potential issues. 🚨\n\n" + ) else: status_msg = "✅ **All Clear:** This PR is ready to merge! 👍\n\n" return markdown_title + status_msg + markdown_output diff --git a/kaizen/reviewer/work_summarizer.py b/kaizen/reviewer/work_summarizer.py index cb59db63..f70e4b9a 100644 --- a/kaizen/reviewer/work_summarizer.py +++ b/kaizen/reviewer/work_summarizer.py @@ -19,11 +19,14 @@ def generate_work_summaries( ): available_tokens = self.provider.available_tokens(WORK_SUMMARY_PROMPT) summaries = [] + # Try to merge the files untill LLM can process the response combined_diff_data = "" total_usage = None for file_dict in diff_file_data: temp_prompt = combined_diff_data temp_prompt += f"""\n---->\nFile Name: {file_dict["file"]}\nPatch: {file_dict["patch"]}\n Status: {file_dict["status"]}""" + + # If available tokens is greated than the new prompt size, process it. if available_tokens - self.provider.get_token_count(temp_prompt) > 0: combined_diff_data = temp_prompt continue @@ -31,6 +34,14 @@ def generate_work_summaries( # Process the prompt prompt = WORK_SUMMARY_PROMPT.format(PATCH_DATA=combined_diff_data) response, usage = self.provider.chat_completion(prompt, user=user) + total_usage = self.provider.update_usage(total_usage, usage) + summaries.append(response) + combined_diff_data = "" + + if combined_diff_data != "": + # process the remaining file diff pending + prompt = WORK_SUMMARY_PROMPT.format(PATCH_DATA=combined_diff_data) + response, usage = self.provider.chat_completion(prompt, user=user) summaries.append(response) combined_diff_data = "" total_usage = self.provider.update_usage(total_usage, usage)