Function to get custom callback for LLM token usage statistics.
(llm_model_name: str)
| 148 | |
| 149 | @contextmanager |
| 150 | def get_custom_callback(llm_model_name: str): |
| 151 | """ |
| 152 | Function to get custom callback for LLM token usage statistics. |
| 153 | """ |
| 154 | cb = CustomCallbackHandler(llm_model_name) |
| 155 | custom_callback.set(cb) |
| 156 | yield cb |
| 157 | custom_callback.set(None) |
no test coverage detected