| 30 | |
| 31 | |
| 32 | class CoT: |
| 33 | def __init__(self, fewshot="\n", model_name="text-davinci-003"): |
| 34 | self.fewshot = fewshot |
| 35 | self.model_name = model_name |
| 36 | self.llm = LLMNode("CoT", model_name, input_type=str, output_type=str) |
| 37 | self.context_prompt = "Answer following questions. Let's think step by step. Give your reasoning process, and then answer the " \ |
| 38 | "question in a new line directly with no extra words.\n" |
| 39 | self.token_unit_price = get_token_unit_price(model_name) |
| 40 | |
| 41 | def run(self, input): |
| 42 | result = {} |
| 43 | st = time.time() |
| 44 | prompt = self.context_prompt + self.fewshot + input + '\n' |
| 45 | response = self.llm.run(prompt, log=True) |
| 46 | result["wall_time"] = time.time() - st |
| 47 | result["input"] = response["input"] |
| 48 | result["output"] = response["output"] |
| 49 | result["prompt_tokens"] = response["prompt_tokens"] |
| 50 | result["completion_tokens"] = response["completion_tokens"] |
| 51 | result["total_tokens"] = response["prompt_tokens"] + response["completion_tokens"] |
| 52 | result["token_cost"] = result["total_tokens"] * self.token_unit_price |
| 53 | result["tool_cost"] = 0 |
| 54 | result["total_cost"] = result["token_cost"] + result["tool_cost"] |
| 55 | result["steps"] = response["output"].count("Step") |
| 56 | return result |
| 57 | |