| 68 | |
| 69 | @dataclass |
| 70 | class TokenTracker: |
| 71 | current_stage: str = "unscoped" |
| 72 | embedding_prompt_tokens: int = 0 |
| 73 | embedding_total_tokens: int = 0 |
| 74 | llm_prompt_tokens: int = 0 |
| 75 | llm_completion_tokens: int = 0 |
| 76 | llm_total_tokens: int = 0 |
| 77 | embedding_requests: int = 0 |
| 78 | llm_requests: int = 0 |
| 79 | rows: List[Dict[str, Any]] = field(default_factory=list) |
| 80 | |
| 81 | @contextmanager |
| 82 | def stage(self, name: str): |
| 83 | previous = self.current_stage |
| 84 | self.current_stage = name |
| 85 | try: |
| 86 | yield |
| 87 | finally: |
| 88 | self.current_stage = previous |
| 89 | |
| 90 | def record_embedding( |
| 91 | self, |
| 92 | *, |
| 93 | model: str, |
| 94 | input_count: int, |
| 95 | usage: Dict[str, int], |
| 96 | provider: str, |
| 97 | ) -> None: |
| 98 | prompt_tokens = int(usage.get("prompt_tokens") or 0) |
| 99 | total_tokens = int(usage.get("total_tokens") or prompt_tokens) |
| 100 | self.embedding_requests += 1 |
| 101 | self.embedding_prompt_tokens += prompt_tokens |
| 102 | self.embedding_total_tokens += total_tokens |
| 103 | self.rows.append( |
| 104 | { |
| 105 | "kind": "embedding", |
| 106 | "stage": self.current_stage, |
| 107 | "provider": provider, |
| 108 | "model": model, |
| 109 | "input_count": input_count, |
| 110 | "prompt_tokens": prompt_tokens, |
| 111 | "completion_tokens": 0, |
| 112 | "total_tokens": total_tokens, |
| 113 | } |
| 114 | ) |
| 115 | |
| 116 | def record_llm( |
| 117 | self, |
| 118 | *, |
| 119 | model: str, |
| 120 | usage: Dict[str, int], |
| 121 | provider: str, |
| 122 | ) -> None: |
| 123 | prompt_tokens = int(usage.get("prompt_tokens") or 0) |
| 124 | completion_tokens = int(usage.get("completion_tokens") or 0) |
| 125 | total_tokens = int(usage.get("total_tokens") or (prompt_tokens + completion_tokens)) |
| 126 | self.llm_requests += 1 |
| 127 | self.llm_prompt_tokens += prompt_tokens |