Run your DAG

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.

example_dag.py

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.

example_dag.py

from trellis_dag import LLM, DAG
from cat_fact_tool import CatFactsAPITool

import asyncio


async def main():
    # Step 2
    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)

    ## Step 3
    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}}}

    ## Step 4
    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 :)