AI: dynamic task scheduling in the context of advanced chatbots — a case study

Photo by José Martín Ramírez Carrasco on Unsplash

If you already read some of my previous articles you already know that I am working on:

an assistant like advanced platform …

AI powered assistant: going further than RAG chatbot

… relying both on directed acyclic graphs (DAGs) …

Directed Acyclic graphs in computer sciences

… and retrieval-augmented generation (RAG).

Retrieval Augmented Generation, or the art of feeding the prompt

This is a personal project that is NOT affiliated with any company.

What I will talk about here is how I use task branching and dynamic scheduling to handle a chat user input processing as a workflow.

We will study the case of a simple workflow with two agents.

The default chat handling directed acyclic graph

1) Agent branching — agent router

As I mentioned in my article about advanced AI assistant, I believe an assistant platform should be multi-agent.

The first task for handling the user input is to determine if it is intended to a particular agent or if it should fallback to a default one.

In the case study example, two agents are available for the Conversational User Interface (CUI), the system agent and the rag agent (the default one)

Starting the input with @system will trigger the “System agent intention router” task while starting the input with @rag (or no agent annotation) will schedule the “RAG agent” task.

@system I want you to learn about black holes and quasars

> will route the query through the system agent

What are black holes?

> no agent annotation given, the rag agent will handle the query

RAG handled query

NB: the agent router will remove the @agent annotation from the input given to the next tasks.

2) Intention branching — system agent intention router

When the agent is determined, we might need to determine the intention conveyed from the input.

In the case study example, the rag agent does not need to use intention branching while the system one does.

What I did in my case was use a zero shot classification to determine the most likely task to use according to the input.

Each of those available tasks have a description, those descriptions and the user input are all converted into embeddings (vector representation of the texte) and by finding the description those embedding is the most similar to the input embedding, we can determine the most likely to be appropriate task to use after.

@system I want to see the current dag

> will match the “I want to see the DAG” description of the task “DAG”

@system I want you to learn about black holes and quasars

> will match the “I want you to learn about a subject” description of the task “Keywords extractor”

“@system I want to see the current dag”

3 ) the input value list branching — forking tasks

One of my tasks is to ingest articles from the arxiv platform. I need to give the API a keyword to use for searching subject related articles.

But what about input asking to ingest multiple keywords?

My first approach was to let the ingestion task work on a list and inside it iterate on each keyword of the list to load articles for a given subject.

This did work but let room for improvements. I then found out about this article “An LLM Compiler for Parallel Function Calling” (december 2023) and used it for inspiration to improve my process.

The new approach is to add a work_index parameter at task definition to let it know which element of the list it should work on.

• If the work_index value is -1 (the default value) and the number of elements in the keyword list is less than two, nothing change.
• If the work_index value is -1 but the keyword list have n elements (n>1), the task will “fork” itself n times, each time with a different work_index value. Each task fork will become a “child” task to the first ingestion one.
• If the work_index value is >= 0, the ingestion task will work on the keyword at the given work_index position

“@system I want you to learn about black holes and quasars”

Using this index based mechanism allows me:

• not to change the task data contract (the required input and provided outputs values),
• parallel ingestion of multiple keywords.

NB: as for now the forked tasks all points to the same task. In the future it might be interesting to enable to define sub-DAG as tasks in case we need multiple tasks to be run one after another for each extracted keyword.

Thanks for reading!

Feel three to follow me or to comment/contact me if you are interested or have constructive criticism to offer!

I feel like the work on this assistant platform thing only have started and I hope you will find it as interesting as I do.