| 11 | |
| 12 | |
| 13 | class ReactBase: |
| 14 | def __init__(self, fewshot, model_name="text-davinci-002", max_iter=8, verbose=True): |
| 15 | self.model_name = model_name |
| 16 | self.max_iter = max_iter |
| 17 | self.verbose = verbose |
| 18 | self.fewshot = fewshot |
| 19 | self.tools = self._load_tools() |
| 20 | if model_name in OPENAI_COMPLETION_MODELS: |
| 21 | self.agent = initialize_agent(self.tools, |
| 22 | OpenAI(temperature=0, model_name=self.model_name), |
| 23 | agent=AgentType.REACT_DOCSTORE, |
| 24 | verbose=self.verbose, |
| 25 | return_intermediate_steps=True, |
| 26 | max_iterations=max_iter) |
| 27 | elif model_name in OPENAI_CHAT_MODELS: |
| 28 | self.agent = initialize_agent(self.tools, |
| 29 | ChatOpenAI(temperature=0, model_name=self.model_name), |
| 30 | agent=AgentType.REACT_DOCSTORE, |
| 31 | verbose=self.verbose, |
| 32 | return_intermediate_steps=True, |
| 33 | max_iterations=max_iter) |
| 34 | self.agent.agent.llm_chain.prompt.template = fewshot |
| 35 | |
| 36 | def run(self, prompt): |
| 37 | self.reset() |
| 38 | result = {} |
| 39 | with get_openai_callback() as cb: |
| 40 | st = time.time() |
| 41 | response = self.agent(prompt) |
| 42 | result["wall_time"] = time.time() - st |
| 43 | result["input"] = response["input"] |
| 44 | result["output"] = response["output"] |
| 45 | result["intermediate_steps"] = response["intermediate_steps"] |
| 46 | result["tool_usage"] = self._parse_tool(response["intermediate_steps"]) |
| 47 | result["total_tokens"] = cb.total_tokens |
| 48 | result["prompt_tokens"] = cb.prompt_tokens |
| 49 | result["completion_tokens"] = cb.completion_tokens |
| 50 | result["total_cost"] = cb.total_cost |
| 51 | result["steps"] = len(response["intermediate_steps"]) + 1 |
| 52 | result["token_cost"] = result["total_cost"] |
| 53 | result["tool_cost"] = 0 |
| 54 | return result |
| 55 | |
| 56 | def _load_tools(self): |
| 57 | docstore = CustomDocstoreExplorer(Wikipedia()) |
| 58 | return [ |
| 59 | Tool( |
| 60 | name="Search", |
| 61 | func=docstore.search, |
| 62 | description="useful for when you need to ask with search" |
| 63 | ), |
| 64 | Tool( |
| 65 | name="Lookup", |
| 66 | func=docstore.lookup, |
| 67 | description="useful for when you need to ask with lookup" |
| 68 | ) |
| 69 | ] |
| 70 | |