| 97 | |
| 98 | |
| 99 | class ReactExtraTool(ReactBase): |
| 100 | def __init__(self, model_name="text-davinci-003", available_tools=["Google", "Calculator"], fewshot="\n", |
| 101 | verbose=True): |
| 102 | self.model_name = model_name |
| 103 | self.verbose = verbose |
| 104 | self.fewshot = fewshot |
| 105 | self.available_tools = available_tools |
| 106 | self.tools = self._load_tools() |
| 107 | self.agent = initialize_agent(self.tools, |
| 108 | OpenAI(temperature=0, model_name=self.model_name), |
| 109 | agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, |
| 110 | verbose=self.verbose, |
| 111 | return_intermediate_steps=True) |
| 112 | |
| 113 | def run(self, prompt): |
| 114 | self.reset() |
| 115 | result = {} |
| 116 | with get_openai_callback() as cb: |
| 117 | st = time.time() |
| 118 | response = self.agent(prompt) |
| 119 | result["wall_time"] = time.time() - st |
| 120 | result["input"] = response["input"] |
| 121 | result["output"] = response["output"] |
| 122 | result["intermediate_steps"] = response["intermediate_steps"] |
| 123 | result["tool_usage"] = self._parse_tool(response["intermediate_steps"]) |
| 124 | result["total_tokens"] = cb.total_tokens + result["tool_usage"]["llm-math_token"] |
| 125 | result["prompt_tokens"] = cb.prompt_tokens |
| 126 | result["completion_tokens"] = cb.completion_tokens |
| 127 | result["total_cost"] = cb.total_cost + result["tool_usage"]["llm-math_token"] * 0.000002 + \ |
| 128 | result["tool_usage"]["serpapi"] * 0.01 # Developer Plan |
| 129 | result["steps"] = len(response["intermediate_steps"]) + 1 |
| 130 | result["token_cost"] = result["total_cost"] |
| 131 | result["tool_cost"] = 0 |
| 132 | |
| 133 | return result |
| 134 | |
| 135 | def _load_tools(self): |
| 136 | tools = [] |
| 137 | for tool_name in self.available_tools: |
| 138 | tool_cls = WORKER_REGISTRY[tool_name] |
| 139 | tools += [Tool(name=tool_name, |
| 140 | func=tool_cls.run, |
| 141 | description=tool_cls.description)] |
| 142 | return tools |
| 143 | |
| 144 | def reset(self): |
| 145 | self.tools = self._load_tools() |
| 146 | self.agent = initialize_agent(self.tools, |
| 147 | OpenAI(temperature=0, model_name=self.model_name), |
| 148 | agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, |
| 149 | verbose=self.verbose, |
| 150 | return_intermediate_steps=True) |
| 151 | self.agent.agent.llm_chain.prompt.template = PREFIX + self._generate_tool_prompt() + "\n" + self.fewshot |
| 152 | |
| 153 | def _parse_tool(self, intermediate_steps): |
| 154 | tool_usage = {"serpapi": 0, "llm-math_token": 0} |
| 155 | for step in intermediate_steps: |
| 156 | if step[0].tool == "Search": |