langfun 0.1.2.dev202512280805__tar.gz → 0.1.2.dev202601160805__tar.gz

{langfun-0.1.2.dev202512280805 → langfun-0.1.2.dev202601160805}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202512280805
+Version: 0.1.2.dev202601160805
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202512280805 → langfun-0.1.2.dev202601160805}/langfun/core/language_model.py

@@ -212,15 +212,30 @@ class ModelInfo(pg.Object):
         The estimated cost in US dollars. If None, cost estimating is not
         supported on the model.
       """
-      # NOTE(daiyip): supported cached tokens accounting in future.
+
       if (self.cost_per_1m_input_tokens is None
           or self.cost_per_1m_output_tokens is None):
         return None
-      return (
-          self.cost_per_1m_input_tokens * usage.prompt_tokens
-          + self.cost_per_1m_output_tokens
-          * (usage.total_tokens - usage.prompt_tokens)
-      ) / 1000_000
+
+      # Calculate cost for non-cached prompt tokens
+      non_cached_prompt_tokens = (
+          usage.prompt_tokens - usage.cached_prompt_tokens
+      )
+      cost = self.cost_per_1m_input_tokens * non_cached_prompt_tokens
+
+      # Add cost for cached prompt tokens (if pricing is available)
+      if self.cost_per_1m_cached_input_tokens is not None:
+        cost += (
+            self.cost_per_1m_cached_input_tokens * usage.cached_prompt_tokens
+        )
+      else:
+        # Fall back to regular input token pricing for cached tokens
+        cost += self.cost_per_1m_input_tokens * usage.cached_prompt_tokens
+
+      # Add cost for output/completion tokens
+      cost += self.cost_per_1m_output_tokens * usage.completion_tokens
+
+      return cost / 1000_000
 
   pricing: Annotated[
       Pricing | None,
@@ -405,6 +420,7 @@ class LMSamplingUsage(pg.Object):
   prompt_tokens: int
   completion_tokens: int
   total_tokens: int
+  cached_prompt_tokens: int = 0
   num_requests: int = 1
   estimated_cost: Annotated[
       float | None,
@@ -429,6 +445,11 @@ class LMSamplingUsage(pg.Object):
     """Returns the average completion tokens per request."""
     return self.completion_tokens // self.num_requests
 
+  @property
+  def average_cached_prompt_tokens(self) -> int:
+    """Returns the average cached prompt tokens per request."""
+    return self.cached_prompt_tokens // self.num_requests
+
   @property
   def average_total_tokens(self) -> int:
     """Returns the average total tokens per request."""
@@ -453,6 +474,8 @@ class LMSamplingUsage(pg.Object):
     return LMSamplingUsage(
         prompt_tokens=self.prompt_tokens + other.prompt_tokens,
         completion_tokens=self.completion_tokens + other.completion_tokens,
+        cached_prompt_tokens=self.cached_prompt_tokens
+        + other.cached_prompt_tokens,
         total_tokens=self.total_tokens + other.total_tokens,
         num_requests=self.num_requests + other.num_requests,
         estimated_cost=estimated_cost,
@@ -1146,6 +1169,7 @@ class LanguageModel(component.Component):
                 prompt_tokens=usage.prompt_tokens // n,
                 completion_tokens=usage.completion_tokens // n,
                 total_tokens=usage.total_tokens // n,
+                cached_prompt_tokens=usage.cached_prompt_tokens // n,
                 estimated_cost=(
                     usage.estimated_cost / n if usage.estimated_cost else None
                 ),
@@ -1637,7 +1661,14 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
   class AggregatedUsage(pg.Object):
     """Aggregated usage."""
 
-    total: LMSamplingUsage = LMSamplingUsage(0, 0, 0, 0, 0.0)
+    total: LMSamplingUsage = LMSamplingUsage(
+        prompt_tokens=0,
+        completion_tokens=0,
+        total_tokens=0,
+        cached_prompt_tokens=0,
+        num_requests=0,
+        estimated_cost=0.0,
+    )
     breakdown: dict[str, LMSamplingUsage] = {}
 
     def __bool__(self) -> bool: