Use cases

What types of problems are classification problems?

  • Deciding which tools an AI agent should use
  • Sentiment analysis
  • Labeling emails
  • Spam detection
  • Customer Intent detection

In this tutorial we’ll write a classifier that will categorize a customer message into the following categories:

  • refund
  • cancel-order
  • technical-support
  • account-issue
  • question

Pre-requisites

Ensure you have read the first guide: Level 1: Writing AI functions. We will re-use the same client we declared there.

The full code for these tutorials are found here

1. Define AI functions and models in BAML files

To solve this problem using LLMs we want a function of this signature: ClassifyMessage(string) -> Category

To do this we can use an enum to represent all the labels we want to classify.

classifier.baml
enum Category {
    Refund
    CancelOrder
    TechnicalSupport
    AccountIssue
    Question
}

Then declare a BAML function signature:

classifier.baml
function ClassifyMessage { 
  input string
  output Category
}

This is more like an interface than an actual function. See next step for how to implement it.

2. Implement the function using an LLM

1

Define the LLM client

2

Define a prompt

Create the prompt by implementing the function using an LLM. In BAML we provide helper utilities to inject the input variables into the prompt, and also get the LLM to return the right output type. You always get full-view of the whole prompt string, without any magic.

impl<llm, ClassifyMessage> version1 {
  client GPT4
  prompt #"
    Classify the following INPUT into ONE
    of the following Intents:

    - Refund
    - CancelOrder
    - TechnicalSupport
    - AccountIssue
    - Question

    INPUT: {#input}

    Response:
  "#
}

2.1 Simplify the prompt using the print_enum utility

Rather than writing out the enums painstakingly into the prompt manually, lets automate that process so that if we change the enum names or values the prompt gets updated automatically.

BAML provides a utility function called print_enum(..) that will print an enum’s values into your prompt in a neatly formatted way that LLMs should understand.

Here is the updated impl:

impl<llm, ClassifyMessage> level1 {
  client GPT4
  prompt #"
    Classify the following INPUT into ONE
    of the following Categories:
    {#print_enum(Category)}

    INPUT: {#input}

    Respond only with the name / identifier. Not any other description.
    Category:
  "#
}

Our VSCode playground shows you what print_enum(..) does to the prompt at compile time. To open it, click on “Open Playground” above the impl.

Here is what the playground will show:

Classify the following INPUT into ONE of the following Categories:

Category
---
Refund
CancelOrder
TechnicalSupport
AccountIssue
Question

INPUT: {arg}

Respond only with the name / identifier. Not any other description.
Category:

You can see it adds the enum name “Category”, and then the list of enums below.

If you need to change the formatting, contact us to learn more about your usecase.

When using print_enum it automatically handles scenarios where you add new enum values, or change the enum name. No need to edit both the prompt and the enum definition.

3. Use the function in your Python / TS code

Our VSCode extension automatically generates a python baml_client to access and call your functions.

from baml_client import baml as b
# The Category type is generated from your BAML code.
from baml_client.baml_types import Category
import asyncio


async def main():
  category = await b.ClassifyMessage("I want to cancel my order")
  if category == Category.CancelOrder:
    print("Customer wants to cancel order")
  else:
    print("Customer wants to {}".format(category))


if __name__ == "__main__":
  asyncio.run(main())

As you keep working with BAML, we recommend setting up a test suite or using our tracing capabilities to make it easier to dive into issues (and get some nice analytics in Boundary Studio).

Conclusion

In this tutorial we’ve learned:

  1. How to use enums to represent a classification problem
  2. BAML utility function — print_enum(..) — to automatically print the enum values into the prompt

Next Steps

In the next section you’ll learn how to use @alias and @description to dynamically change the name of the enums, while keeping your python code intact.