| 12 | self.token_unit_price = get_token_unit_price(model_name) |
| 13 | |
| 14 | def run(self, input): |
| 15 | result = {} |
| 16 | st = time.time() |
| 17 | prompt = self.context_prompt + self.fewshot + input + '\n' |
| 18 | response = self.llm.run(prompt, log=True) |
| 19 | result["wall_time"] = time.time() - st |
| 20 | result["input"] = response["input"] |
| 21 | result["output"] = response["output"] |
| 22 | result["prompt_tokens"] = response["prompt_tokens"] |
| 23 | result["completion_tokens"] = response["completion_tokens"] |
| 24 | result["total_tokens"] = response["prompt_tokens"] + response["completion_tokens"] |
| 25 | result["token_cost"] = result["total_tokens"] * self.token_unit_price |
| 26 | result["tool_cost"] = 0 |
| 27 | result["total_cost"] = result["token_cost"] + result["tool_cost"] |
| 28 | result["steps"] = 1 |
| 29 | return result |
| 30 | |
| 31 | |
| 32 | class CoT: |