diff --git a/src/main/java/com/llmproxy/service/llm/TokenEstimator.java b/src/main/java/com/llmproxy/service/llm/TokenEstimator.java
index 1963501..8a97bd6 100644
--- a/src/main/java/com/llmproxy/service/llm/TokenEstimator.java
+++ b/src/main/java/com/llmproxy/service/llm/TokenEstimator.java
@@ -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;
@@ -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
         }
     }
-}
+}
\ No newline at end of file