import time
from langchain.agents import AgentType
from langchain.agents import initialize_agent
from langchain.callbacks import get_openai_callback
from langchain.chat_models import ChatOpenAI
from nodes.Worker import *
from utils.CustomDocstoreExplorer import CustomDocstoreExplorer
from utils.util import *
class ReactBase:
def __init__(self, fewshot, model_name="text-davinci-002", max_iter=8, verbose=True):
self.model_name = model_name
self.max_iter = max_iter
self.verbose = verbose
self.fewshot = fewshot
self.tools = self._load_tools()
if model_name in OPENAI_COMPLETION_MODELS:
self.agent = initialize_agent(self.tools,
OpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.REACT_DOCSTORE,
verbose=self.verbose,
return_intermediate_steps=True,
max_iterations=max_iter)
elif model_name in OPENAI_CHAT_MODELS:
self.agent = initialize_agent(self.tools,
ChatOpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.REACT_DOCSTORE,
verbose=self.verbose,
return_intermediate_steps=True,
max_iterations=max_iter)
self.agent.agent.llm_chain.prompt.template = fewshot
def run(self, prompt):
self.reset()
result = {}
with get_openai_callback() as cb:
st = time.time()
response = self.agent(prompt)
result["wall_time"] = time.time() - st
result["input"] = response["input"]
result["output"] = response["output"]
result["intermediate_steps"] = response["intermediate_steps"]
result["tool_usage"] = self._parse_tool(response["intermediate_steps"])
result["total_tokens"] = cb.total_tokens
result["prompt_tokens"] = cb.prompt_tokens
result["completion_tokens"] = cb.completion_tokens
result["total_cost"] = cb.total_cost
result["steps"] = len(response["intermediate_steps"]) + 1
result["token_cost"] = result["total_cost"]
result["tool_cost"] = 0
return result
def _load_tools(self):
docstore = CustomDocstoreExplorer(Wikipedia())
return [
Tool(
name="Search",
func=docstore.search,
description="useful for when you need to ask with search"
),
Tool(
name="Lookup",
func=docstore.lookup,
description="useful for when you need to ask with lookup"
)
]
def reset(self):
self.tools = self._load_tools()
if self.model_name in OPENAI_COMPLETION_MODELS:
self.agent = initialize_agent(self.tools,
OpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.REACT_DOCSTORE,
verbose=self.verbose,
return_intermediate_steps=True,
max_iterations=self.max_iter)
elif self.model_name in OPENAI_CHAT_MODELS:
self.agent = initialize_agent(self.tools,
ChatOpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.REACT_DOCSTORE,
verbose=self.verbose,
return_intermediate_steps=True,
max_iterations=self.max_iter)
self.agent.agent.llm_chain.prompt.template = self.fewshot
def _parse_tool(self, intermediate_steps):
tool_usage = {"search": 0, "lookup": 0}
for step in intermediate_steps:
if step[0].tool == "Search":
tool_usage["search"] += 1
if step[0].tool == "Lookup":
tool_usage["lookup"] += 1
return tool_usage
class ReactExtraTool(ReactBase):
def __init__(self, model_name="text-davinci-003", available_tools=["Google", "Calculator"], fewshot="\n",
verbose=True):
self.model_name = model_name
self.verbose = verbose
self.fewshot = fewshot
self.available_tools = available_tools
self.tools = self._load_tools()
self.agent = initialize_agent(self.tools,
OpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
verbose=self.verbose,
return_intermediate_steps=True)
def run(self, prompt):
self.reset()
result = {}
with get_openai_callback() as cb:
st = time.time()
response = self.agent(prompt)
result["wall_time"] = time.time() - st
result["input"] = response["input"]
result["output"] = response["output"]
result["intermediate_steps"] = response["intermediate_steps"]
result["tool_usage"] = self._parse_tool(response["intermediate_steps"])
result["total_tokens"] = cb.total_tokens + result["tool_usage"]["llm-math_token"]
result["prompt_tokens"] = cb.prompt_tokens
result["completion_tokens"] = cb.completion_tokens
result["total_cost"] = cb.total_cost + result["tool_usage"]["llm-math_token"] * 0.000002 + \
result["tool_usage"]["serpapi"] * 0.01 # Developer Plan
result["steps"] = len(response["intermediate_steps"]) + 1
result["token_cost"] = result["total_cost"]
result["tool_cost"] = 0
return result
def _load_tools(self):
tools = []
for tool_name in self.available_tools:
tool_cls = WORKER_REGISTRY[tool_name]
tools += [Tool(name=tool_name,
func=tool_cls.run,
description=tool_cls.description)]
return tools
def reset(self):
self.tools = self._load_tools()
self.agent = initialize_agent(self.tools,
OpenAI(temperature=0, model_name=self.model_name),
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
verbose=self.verbose,
return_intermediate_steps=True)
self.agent.agent.llm_chain.prompt.template = PREFIX + self._generate_tool_prompt() + "\n" + self.fewshot
def _parse_tool(self, intermediate_steps):
tool_usage = {"serpapi": 0, "llm-math_token": 0}
for step in intermediate_steps:
if step[0].tool == "Search":
tool_usage["serpapi"] += 1
if step[0].tool == "Calculator":
tool_usage["llm-math_token"] += len(step[0].tool_input + step[1]) // 4 # 4 chars per token
return tool_usage
def _get_worker(self, name):
if name in WORKER_REGISTRY:
return WORKER_REGISTRY[name]
else:
raise ValueError("Worker not found")
def _generate_tool_prompt(self):
prompt = "Tools can be one of the following:\n"
for name in self.available_tools:
worker = self._get_worker(name)
prompt += f"{worker.name}[input]: {worker.description}\n"
return prompt + "\n"
PREFIX = """Answer the following questions as best you can. You have access to the following tools:
"""