omkarenator's picture
Make ActionRunner accept custom LLMs (#13)
f893655 unverified
raw
history blame contribute delete
No virus
2.27 kB
import os
from duckpy import Client
from langchain import PromptTemplate, OpenAI, LLMChain
from langchain.agents import Tool
from langchain.base_language import BaseLanguageModel
MAX_SEARCH_RESULTS = 20 # Number of search results to observe at a time
search_description = """ Useful for when you need to ask with search. Use direct language and be
EXPLICIT in what you want to search.
## Examples of incorrect use
1. Action: Search
Action Input: "[name of bagel shop] menu"
The Action Input cannot be None or empty.
"""
notepad_description = """ Useful for when you need to note-down specific
information for later reference. Please provide full information you want to
note-down in the Action Input and all future prompts will remember it.
This is the mandatory tool after using the search tool.
Using Notepad does not always lead to a final answer.
## Exampels of using notepad tool
Action: Notepad
Action Input: the information you want to note-down
"""
async def ddg(query: str):
if query is None or query.lower().strip().strip('"') == "none" or query.lower().strip().strip('"') == "null":
x = "The action input field is empty. Please provide a search query."
return [x]
else:
client = Client()
return client.search(query)[:MAX_SEARCH_RESULTS]
async def notepad(x: str) -> str:
return f"{[x]}"
search_tool = Tool(name="Search",
func=lambda x: x,
coroutine=ddg,
description=search_description)
note_tool = Tool(name="Notepad",
func=lambda x: x,
coroutine=notepad,
description=notepad_description)
def rewrite_search_query(q: str, search_history, llm: BaseLanguageModel) -> str:
history_string = '\n'.join(search_history)
template ="""We are using the Search tool.
# Previous queries:
{history_string}. \n\n Rewrite query {action_input} to be
different from the previous ones."""
prompt = PromptTemplate(template=template,
input_variables=["action_input", "history_string"])
llm_chain = LLMChain(prompt=prompt, llm=llm)
return llm_chain.predict(action_input=q, history_string=history_string)