Tutorial Overview
The final step in our tutorial is to run our DAG. We’ll need to set one argument for our dag’s execute
function, and move everything into an async function.
Setting the execute args
If you will remember, we didn’t set the arguments for our CatFactTool
, so we will now do so with set_execute_args
. The function takes a dictionary of the format: {"args": list, "kwargs": dict}
. In our case, we have max_length = 100
and limit = 1
.
init_dict = {cat_fact_tool.get_id(): {"kwargs": {"max_length": 100}}}
result = await dag.execute(init_dict)
Running with asyncio
The last step is to move all the code into an async function called main, and run it with asyncio.
from trellis_dag import LLM, DAG
from cat_fact_tool import CatFactsAPITool
import asyncio
async def main():
generate_cat_fact_llm_msgs = [
{
"role": "user",
"content": "Tell me a random cat fact, as a sentence.",
}
]
generate_cat_fact_llm = LLM(
"generate_cat_fact_llm", messages=generate_cat_fact_llm_msgs
)
distinguish_cat_fact_llm_msgs = [
{
"role": "user",
"content": "Which of these was generated by an LLM? 1. {cat_fact_1} 2. {cat_fact_2} Give your answer as 1 or 2.",
}
]
distinguish_cat_fact_llm = LLM("distinguish_cat_fact_llm")
distinguish_cat_fact_llm.set_messages(distinguish_cat_fact_llm_msgs)
cat_fact_tool = CatFactsAPITool("cat_fact_tool")
dag = DAG()
dag.add_node(cat_fact_tool)
dag.add_node(generate_cat_fact_llm)
dag.add_node(distinguish_cat_fact_llm)
dag.add_edge(cat_fact_tool, distinguish_cat_fact_llm)
dag.add_edge(
generate_cat_fact_llm,
distinguish_cat_fact_llm,
fn=lambda x: {"cat_fact_2": x["choices"][0]["message"]["content"]},
)
init_dict = {cat_fact_tool.get_id(): {"kwargs": {"max_length": 100}}}
result = await dag.execute(init_dict)
print(result)
if __name__ == "__main__":
asyncio.run(main())
And with that… this tutorial is finished! We’re excited to see what you build with Trellis :)