Use cases

What types of problems are entity extraction problems?

  • Filling out a form with a chat bot
  • Getting specific details out of long form text
  • Calling a function in your code with some parameters after analyzing some text (like OpenAI Functions)

Pre-requisites

Extracting order details from text

In this tutorial we’ll write an AI function that can take an Order status email and extract order information from it.

Our function signature will be: GetOrderInfo(email: Email) -> OrderInfo

Here is the reference code for these sets of tutorials. Clone the repo and install BAML extension to get syntax highlighting!

1. Define AI functions and models in BAML files

First create the data models using classes, which actually get converted to actual pydantic models at generation time if using Python.

baml_src/main.baml
class Email {
  subject string
  body string
}

class OrderInfo {
  id string
  date string
  product_name string
  cost float
}

The code snippet you provided is declaring the AI function signature for a function called GetOrderInfo. This function takes an input of type Email and outputs an OrderInfo. By defining this function signature in your BAML file, you are essentially setting up the structure for how the AI function will operate and what kind of data it will expect as input and output. Then declare the AI function signature

baml_src/main.baml
function GetOrderInfo {
  input Email
  output OrderInfo
}

Hit Ctrl + S in VSCode, and you should get a generated baml_client to import in your python code (more on this later).

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.

impl<llm, GetOrderInfo> version1 {
  client GPT4
  prompt #"
    Given the email below:

    Email Subject: {#input.subject}
    Email Body: {#input.body}

    Extract this info from the email in JSON format:
    {
      "id": string,
      "date": string,
      "product_name": string,
      "cost": float
    }

    JSON:
  "#
}

3. Simplify the prompt using print_type utility

Rather than writing out the JSON format we want manually, lets automate that process so that if we add, remove or change fields the prompt is automatically updated.

BAML provides a utility called print_type that will print out the JSON format of a type. We’ll use this to print out the desired OrderInfo type.

...
  prompt #"
    Given the email below:

    Email Subject: {#input.subject}
    Email Body: {#input.body}

    Extract this info from the email in JSON format:
    {#print_type(output)}

    JSON:
  "#
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.

In the BAML VSCode Playground you should see the rendered prompt:

Given the email below:

Email Subject: {arg.subject}
Email Body: {arg.body}

Extract this info from the email in JSON format:
{
  "id": string,
  "date": string,
  "product_name": string,
  "cost": float
}

JSON:

3. Call the function in Python / TS

Our VSCode extension automatically generates a python baml_client to access and call your functions. All fully type-safe. Verify that if you add the wrong input you do see an error.

from baml_client import baml as b
# Import your generated Email model. 
# We generate this pydantic model for you.
from baml_client.baml_types import Email
import asyncio

async def main():
  order_info = await b.GetOrderInfo(Email(
      subject="Order #1234",
      body="Your order has been shipped. It will arrive on 1st Jan 2022. Product: iPhone 13. Cost: $999.99"
  ))

  # This is all fully typed for you:
  if order_info.cost > 1000:
      print("You spent a lot of money!")
  elif order_info.cost > 500:
      print("You spent a chunk of money!")
  else:
      print("You spent a little money!")

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

Note that you may get some errors testing other cases. We’ll discuss how to make this more resilient in the next tutorial.

Full code

Check it out!

Conclusion

In this tutorial we learned:

  1. Basic entity extraction tasks
  2. print_type(..) prompt engineering utility
  3. How to call AI functions in Python / TS