coderpotter's picture
Upload folder using huggingface_hub
7b2e5db verified
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from research_assistant.components.agent import Agent
from research_assistant.config.configuration import ConfigurationManager
class PlannerParser:
tool_list_schema = {
"title": "ToolSchema",
"description": "This is the schema used to get tools list, after parsing the plan string, which will be used to do tool or fucntion calling for the further parts of the agentic framework",
"type": "object",
"properties": {
"tools": {
"type": "array",
"items": {"type": "string"},
"title": "Tools",
"description": "The Tool each step in the plan needs to access. If we have seven steps, this has a list of 7 values depicting the tool name for each step that needs to be used",
"default": [],
},
},
"required": ["tools"],
}
argument_list_schema = {
"title": "ArgumentSchema",
"description": "This is the schema used to get argument list used for each tool call, after parsing the plan string, which will be used to do tool or fucntion calling for the further parts of the agentic framework",
"type": "object",
"properties": {
"arguments": {
"type": "array",
"items": {"type": "string"},
"title": "Arguments",
"description": "The Arugment for tool call for each step in the plan needs to access. If we have seven steps, this has a list of 7 values depicting the argument value for each step that needs to be used",
"default": [],
},
},
"required": ["arguments"],
}
dependency_list_schema = {
"title": "DependentSchema",
"description": "This schema defines dependencies for each tool call to fetch additional context after parsing the plan string for the agentic framework.",
"type": "object",
"properties": {
"dependencies": {
"type": "object",
"additionalProperties": {
"type": "array",
"items": {"type": "integer"},
},
"title": "Dependencies",
"description": "A dictionary where each key is a step identifier, and the value is an array of integers, representing the step dependencies.",
"default": {},
}
},
"required": ["dependencies"],
}
def __init__(self, plan_string: str):
self.plan_string = plan_string
self.config = ConfigurationManager().get_planner_parser_config()
def get_tool_list(self):
llm = Agent(self.config.tool_list_model).get_model()
prompt_template = ChatPromptTemplate.from_messages(
[
SystemMessage(
(
"You are a planner parser. You will be given a plan consisting of a series of steps, and you need to give me a list consisting of what tool is being used for each step.\n"
"So if there are 7 steps, I need a list of length 7, where each value is the tool that is being used for the corresponding step.\n"
"For example:\n"
'1. Find out the temperature right now. #E1 = Google["What is the temperature near me?"]\n'
'2. Check the weather forecast for tomorrow. #E2: WeatherAPI["What will be the weather like tomorrow?"]\n'
'3. Get the top 5 news articles related to the topic of climate change. !3 = NewsAPI["What are the top 5 news articles about climate change?"]\n'
"If this is the given plan, the output should be: ['Google', 'WeatherAPI', 'NewsAPI'].\n"
"Be very careful that you don't miss any step or tool. The number of steps and tools should be the same. Check your output thoroughly. Write only one tool on each line."
),
),
MessagesPlaceholder(variable_name="plan"),
MessagesPlaceholder(variable_name="context", optional=True),
]
)
tool_parser = prompt_template | llm.with_structured_output(
self.tool_list_schema
)
return tool_parser.invoke(
{
"plan": [
HumanMessage(
f"This is the generated plan:\n{self.plan_string}\n\n"
"Now parse this content and give me the list of tools for each step."
)
]
}
)["tools"]
def get_argument_list(self):
llm = Agent(self.config.argument_list_model).get_model()
prompt_template = ChatPromptTemplate.from_messages(
[
SystemMessage(
(
"You are a planner parser. You take in a plan consisting of a series of steps, and you need to give me a list consisting of what argument is being used for each tool call in each step.\n"
"So if there are 7 steps, I need a list of length 7, where each value is the argument that is being called inside the tool for each step. A step is defined as #Ex where x is the step number. An argument will always be of the format #Ex = tool_name['argument']\n"
"For example:\n"
'1. Find out the temperature right now. #E1 = Google["What is the temperature near me?"]\n'
'2. Check the weather forecast for tomorrow. #E2 : WeatherAPI["What will be the weather like tomorrow?"]\n'
'3. Get the top 5 news articles related to the topic of climate change. !3 = NewsAPI["What are the top 5 news articles about climate change?"]\n'
"If this is the given plan, the output should be: ['What is the temperature near me?', 'What will be the weather like tomorrow?', 'What are the top 5 news articles about climate change?']\n"
"Be very careful that you don't miss any step or argument. The number of steps and arguments should be the same. Check your output thoroughly."
),
),
MessagesPlaceholder(variable_name="plan"),
MessagesPlaceholder(variable_name="context", optional=True),
]
)
argument_parser = prompt_template | llm.with_structured_output(
self.argument_list_schema
)
return argument_parser.invoke(
{
"plan": [
HumanMessage(
f"This is the generated plan:\n{self.plan_string}\n\n"
"Give me the list of arguments of each tool call for each step. For each step I need to know what is the argument that is being passed inside the tool. If I have 10 steps, I need a list of length 10, consisting of the query or argument for each tool call in each step."
)
]
}
)["arguments"]
def get_dependency_list(self):
llm = Agent(self.config.dependency_list_model).get_model()
prompt_template = ChatPromptTemplate.from_messages(
[
SystemMessage(
(
"You are a planner parser. You get a plan consisting of a series of steps, and you need to give me a dictionary consisting of what step results each step argument is dependent upon. If there are 7 steps and 5 steps require the results of previous steps, I need a dictionary containing those 5 keys, where each value is a list of step numbers that the key step is dependent upon. For example:\n"
'1. Find out the temperature right now. #E1 = Location["What is my current location?"]\n'
'2. Check the weather forecast for tomorrow. #E2 : WeatherAPI["What will be the weather like tomorrow at #E1?"]\n'
'3. Get the top 5 news articles related to the topic of climate change. !3 = LLM["What are major tourist things to do based on the information of #E1 and #E2?"]\n'
"If this is the given plan, the output should be:\n"
"{'2': [1], '3': [1, 2]}\n"
"This is because step 2 is dependent on the value of #E1, and step 3 is dependent on the value of #E2 and #E1.\n"
"A step cannot depend on itself. A step cannot depend on any step that comes after it. A dependency only exists if a step's content has 'from #Ex' where x is the step number.\n"
"For example:\n"
'#E8 = LLM["What were the main results and findings from their experiments? Explain the results in simple terms without technical jargon"]\n'
'#E9 = Arxiv["UNLI limitations natural language inference"]\n'
"Under no circumstances can the output be: {'8': [9]}\n"
"This is because 8 comes before 9.\n\n"
"Be very careful that you don't miss any step or dependency. The number of steps and dependencies should be the same. Check your output thoroughly."
),
),
MessagesPlaceholder(variable_name="plan"),
MessagesPlaceholder(variable_name="context", optional=True),
]
)
dependency_parser = prompt_template | llm.with_structured_output(
self.dependency_list_schema
)
return dependency_parser.invoke(
{
"plan": [
HumanMessage(
f"This is the generated plan:\n{self.plan_string}\n\n"
"Parse this content and give me the dictionary of dependencies. Look at each step and see on what steps each argument of the step is dependent upon, and give me the values in the form of a key value pair. Key being the step number and value being the list of step numbers that the key is dependent upon."
)
]
}
)["dependencies"]