Skip to content

Artemis: feat: Improve Token Estimator Documentation and Logic Clarity #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mike-turintech wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Artemis: feat: Improve Token Estimator Documentation and Logic Clarity #4

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 22 additions & 5 deletions src/main/java/com/llmproxy/service/llm/TokenEstimator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,12 @@

@Component
public class TokenEstimator {
/**
* Estimates the number of tokens in the given text using a simple length-based heuristic.
*
* @param text the text to analyze
* @return the estimated token count (text length divided by 4)
*/
public int estimateTokenCount(String text) {
if (text == null || text.isEmpty()) {
return 0;
Expand All @@ -12,12 +18,23 @@ public int estimateTokenCount(String text) {
return text.length() / 4;
}

/**
* Updates the QueryResult with token count estimates if not already set.
*
* @param result the query result to update
* @param query the input query text
* @param response the response text
*/
public void estimateTokens(QueryResult result, String query, String response) {
if (result.getTotalTokens() == 0) {
result.setInputTokens(estimateTokenCount(query));
result.setOutputTokens(estimateTokenCount(response));
result.setTotalTokens(result.getInputTokens() + result.getOutputTokens());
result.setNumTokens(result.getTotalTokens()); // For backward compatibility
int inputTokens = estimateTokenCount(query);
int outputTokens = estimateTokenCount(response);
int totalTokens = inputTokens + outputTokens;

result.setInputTokens(inputTokens);
result.setOutputTokens(outputTokens);
result.setTotalTokens(totalTokens);
result.setNumTokens(totalTokens); // For backward compatibility
}
}
}
}