DEV Community: Stefan Krawczyk

Building an Email Assistant Application with Burr

Stefan Krawczyk — Fri, 26 Apr 2024 20:46:18 +0000

In this tutorial, I will demonstrate how to use Burr, an open source framework (disclosure: I helped create it), using simple OpenAI client calls to GPT4, and FastAPI to create a custom email assistant agent. We’ll describe the challenge one faces and then how you can solve for them. For the application frontend we provide a reference implementation but won’t dive into details for it.

Why are interactive agents applications a challenge?

LLMs rarely achieve complex goals on their own, and almost never on the first try. While it is in vogue to claim that ChatGPT given an internet connection can solve the world’s problems, the majority of high-value tools we’ve encountered use a blend of AI ingenuity and human guidance. This is part of the general move towards building agents — an approach where the AI makes decisions from information it receives — this could be information it queries, information a user provides, or information another LLM gives it.

A simple example of this is a tool to help you draft a response to an email. You put the email and your response goals, and it writes the response for you. At a minimum, you’ll want to provide feedback so it can adjust the response. Furthermore, you will want it to give a chance to ask clarifying questions (an overly confident yet incorrect chatbot helps no one).

In designing this interaction, your system will, inevitably, become a back-and-forth between user/LLM control. In addition to the standard challenges around AI applications (unreliable APIs, stochastic implementations, etc…), you will face a suite of new problems, including:

Logically modeling a set of interaction points/flows
Persisting the state so the user can pick up the interaction/application from where it left off
Monitoring the decisions the LLM made (E.G. whether to ask the user questions or not)

And so on… In this post we’re going to walk through how to approach solving these — we’ll use the Burr library as well as FastAPI to build a web service to address these challenges in an extensible, modular manner; so you can then use this as a blue print for your own agent assistant needs.

The Tools

Burr

Burr is a lightweight python library you use to build applications as state machines. You construct your application out of a series of actions (these can be either decorated functions or objects), which declare inputs from state, as well as inputs from the user. These specify custom logic (delegating to any framework), as well as instructions on how to update state. State is immutable, which allows you to inspect it at any given point. Burr handles orchestration, monitoring and persistence.



@action(reads=["counter"], writes=["counter"])
def count(state: State) -> Tuple[dict, State]:
    current = state["counter"] + 1
    result = {"counter": current}
    return result, state.update(counter=counter)

Note that the action above has two returns — the results (the counter), and the new, modified state (with the counter field incremented).

You run your Burr actions as part of an application — this allows you to string them together with a series of (optionally) conditional transitions from action to action.



from burr.core import ApplicationBuilder, default, expr
app = (
    ApplicationBuilder()
    .with_state(counter=0) # initialize the count to zero
    .with_actions(
        count=count, 
        done=done # implementation left out above
    ).with_transitions(
        ("count", "count", expr("counter < 10")), # Keep counting if the counter is less than 10
        ("count", "done", default) # Otherwise, we're done
    ).with_entrypoint("count") # we have to start somewhere
    .build()
)

Burr comes with a user-interface that enables monitoring/telemetry, as well as hooks to persist state/execute arbitrary code during execution.

You can visualize this as a flow chart, i.e. graph / state machine:

Image of our application as produced by Burr. Image by author.

And monitor it using the local telemetry debugger:

Burr comes with a UI — this is what it looks like when inspecting a run of our counter example. Image by author.

While we showed the (very simple) counter example above, Burr is more commonly used for building chatbots/agents (we’ll be going over an example in this post).

FastAPI

FastAPI is a framework that lets you expose python functions in a REST API. It has a simple interface — you write your functions then decorate them, and run your script — turning it into a server with self-documenting endpoints through OpenAPI.



@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.get("/items/{item_id}")
def read_item(item_id: int, q: Union[str, None] = None):
    """A very simpler example of an endpoint that takes in arguments."""
    return {"item_id": item_id, "q": q}

FastAPI is easy to deploy on any cloud provider — it is infrastructure-agnostic and can generally scale horizontally (so long as consideration into state management is done). See this page for more information.

React (or any frontend framework)

You can use any frontend framework you want — react-based tooling, however, has a natural advantage as it models everything as a function of state, which can map 1:1 with the concept in Burr. In the demo app we use react, react-query, and tailwind, but we’ll be skipping over this largely (it is not central to the purpose of the post).

Building

Let’s dig a bit more into the conceptual model. At a high-level, our email assistant will do the following:

Accept an email + instructions to respond
Come up with a set of clarifying questions (if the LLM deems it required)
Generates a draft using the answer to those questions
Accept feedback to that draft and generates another one, repeating until the user is happy
Return the final draft (done)

Modeling Control Flow

As Burr requires you to build a control flow from actions and transitions, we can initially model this as a simple flowchart.

What our application will look like. Image by author.

We drafted this before actually writing any code — you will see it transforms to code naturally.

The green nodes represent actions (these take state in and modify it), and the blue nodes represent inputs (these are points at which the app has to pause and ask the user for information). Note that there is a loop (formulate_draft ⇔process_feedback) — we iterate on feedback until we’re happy with the results.

This diagram is simply a stylized version of what Burr shows you — the modeling is meant to be close to the actual code. We have not displayed state information (the data the steps take in/return), but we’ll need to track the following (that may or may not be populated at any given point) so we can make decisions about what to do next:

The initial inputs: {email_to_respond: str, response_instructions: str}
The questions the LLM asks and the user responses (if any):{clarifications: list[str], response_instructions: list[str]}
The list of drafts + feedback: {drafts: list[str], feedback_history: list[str]}
The final result: {final_result: str}

Implementing/Testing

Looking at the requirements above, we can build a straightforward burr application since we can very closely match our code with our diagram above. Let’s take a look at the determine_clarifications step, for example:



@action(
    reads=["response_instructions", "incoming_email"], 
    writes=["clarification_questions"]
)
def determine_clarifications(state: State) -> Tuple[dict, State]:
    """Determines if the response instructions require clarification."""

    incoming_email = state["incoming_email"]
    response_instructions = state["response_instructions"]
    client = _get_openai_client()

    result = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": ("You are a chatbot that has the task of "
                            "generating responses to an email on behalf "
                            "of a user. "),
            },
            {
                "role": "user",
                "content": (
                    f"The email you are to respond to is: {incoming_email}."
                     # ... left out, see link above
                    "The questions, joined by newlines, must be the only "
                    "text you return. If you do not need clarification, "
                    "return an empty string."
                ),
            },
        ],
    )
    content = result.choices[0].message.content
    all_questions = content.split("\n") if content else []
    return {"clarification_questions": all_questions}, state.update(
        clarification_questions=all_questions)

Note that this uses simple OpenAI calls — you can replace this with Langchain, LlamaIndex, Hamilton (or something else) if you prefer more abstraction, and delegate to whatever LLM you like to use. And, you should probably use something a little more concrete (E.G. instructor) to guarantee output shape.

To tie these together, we put them into the application builder — this allows us to set conditional transitions (e.g. len(clarification_questions>0) and therefore connect actions, recreating the diagram above.



application = (
    ApplicationBuilder()
    # define our actions
    .with_actions(
        process_input,
        determine_clarifications,
        clarify_instructions,
        formulate_draft,
        process_feedback,
        final_result,
    )
    # define how our actions connect
    .with_transitions(
        ("process_input", "determine_clarifications"),
        (
            "determine_clarifications",
            "clarify_instructions",
            expr("len(clarification_questions) > 0"),
        ),
        ("determine_clarifications", "formulate_draft"),
        ("clarify_instructions", "formulate_draft"),
        ("formulate_draft", "process_feedback"),
        ("process_feedback", "formulate_draft", expr("len(feedback) > 0")),
        ("process_feedback", "final_result"),
    )
    .with_state(draft_history=[])
    .with_entrypoint("process_input")
    .build()
)

To iterate on this, we used a jupyter notebook. Running our application is simple — all you do is call the .run() method on the Application, with the right halting conditions. We’ll want it to halt before any action that requires user input (clarify_instructions and process_feedback), and after final_result. We can then run it in a while loop, asking for user input and feeding it back to the state machine:



def request_answers(questions):
    """Requests answers from the user for the questions the LLM has"""
    answers = []
    print("The email assistant wants more information:\n")
    for question in questions:
        answers.append(input(question))
    return answers

def request_feedback(draft):
    """Requests feedback from the user for a draft"""
    print( 
        f"here's a draft!: \n {draft} \n \n What feedback do you have?",
    )
    return input("Write feedback or leave blank to continue (if you're happy)")
inputs = {
    "email_to_respond" : EMAIL,
    "response_instructions" : INSTRUCTIONS
}

# in our notebook cell:
while True:
    action, result, state = app.run(
        halt_before=["clarify_instructions", "process_feedback"], 
        halt_after=["final_result"],
        inputs=inputs
    )
    if action.name == "clarify_instructions":
        questions = state["clarification_questions"]
        answers = request_answers(questions)
        inputs = {
            "clarification_inputs" : answers
        }
    if action.name == "process_feedback":
        feedback = request_feedback(state["current_draft"])
        inputs = {"feedback" : feedback}
    if action.name == "final_result":
        print("final result is:", state["current_draft"])
        break

You can then use the Burr UI to monitor your application as it runs!

Example of using the Burr UI (with the email app UI) and then seeing it’s execution. Image by author.

Persistence

We’re going to persist our results to an SQLite server (although as you’ll see later on this is customizable). To do this, we need to add a few lines to the ApplicationBuilder.



state_persister = SQLLitePersister(
    db_path="sqllite.db", 
    table_name="email_assistant_table"
)

app = (
    ApplicationBuilder().
    ... # the code we had above
    .initialize(  
        initializer=state_persister,
        resume_at_next_action=True,
        default_state={"chat_history" : []},
        default_entrypoint="process_input"
    )
    .with_identifiers(app_id=app_id)
    .build()
)

This ensures that every email draft we create will be saved and can be loaded at every step. When you want to resume a prior draft of an email, all you have to do is rerun the code and it will start where it left off.

Integrating in a web server

To expose this in a web server we’ll be using FastAPI to create endpoints and Pydantic to represent types. Before we get into the details, we’ll note that Burr naturally provides an application_id (either generated or specified) for every instance of an application. In this case the application_id would correspond to a particular email draft. This allows us to uniquely access it, query from the db, etc… It also allows for a partition key (E.G. user_id) so you can add additional indexing in your database. We center the API around inputs/outputs.

Endpoints

We will construct the following endpoints:

POST /create: This will create a new application and return the ID
PUT /initialize_draft/{id}/: This calls out to process_input, passing in the email and instructions
PUT /clarify_instructions/{id}: This will give answers back to the LLM
PUT /process_feedback/{id}: This will give feedback back to the LLM
GET /{id}/state: This will return the current state of the application

The GET endpoint allows us to get the current state of the application — this enables the user to reload if they quit the browser/get distracted. Each of these endpoints will return the full state of the application, which can be rendered on the frontend. Furthermore, it will indicate the next API endpoint we call, which allows the UI to render the appropriate form and submit to the right endpoint.

Using FastAPI + Pydantic, this becomes very simple to implement. First, let’s add a utility to get the application object. This will use a cached version or instantiate it:



@functools.lru_cache(maxsize=128)
def get_application(app_id: str) -> Application:
    app = email_assistant_application.application(app_id=app_id)
    return app

All this does is call our function application in email_assistant that recreates the application. We have not included the create function here, but it calls out to the same API.

Data Model
Let’s then define a Pydantic model to represent the state, and the app object in FastAPI:



class EmailAssistantState(pydantic.BaseModel):
    app_id: str
    email_to_respond: Optional[str]
    response_instructions: Optional[str]
    questions: Optional[List[str]]
    answers: Optional[List[str]]
    drafts: List[str]
    feedback_history: List[str]
    final_draft: Optional[str]
    # This stores the next step, which tells the frontend which ones to call
    next_step: Literal[
        "process_input", "clarify_instructions", 
        "process_feedback", None]

    @staticmethod
    def from_app(app: Application):
        # implementation left out, call app.state and translate to 
        # pydantic model we can use `app.get_next_action()` to get 
        #the next step and return it to the user
        ...

Note that every endpoint will return this same pydantic model!

Endpoints

Given that each endpoint returns the same thing (a representation of the current state as well as the next step to execute), they all look the same. We can first implement a generic run_through function, which will progress our state machine forward, and return the state.



def run_through(
    project_id: str, 
    app_id: Optional[str], 
    inputs: Dict[str, Any]
) -> EmailAssistantState:
    email_assistant_app = get_application(project_id, app_id)
    email_assistant_app.run(
        halt_before=["clarify_instructions", "process_feedback"],
        halt_after=["final_result"],
        inputs=inputs,
    )
    return EmailAssistantState.from_app(email_assistant_app)

This represents a simple but powerful architecture. We can continue calling these endpoints until we’re at a “terminal” state, at which point we can always ask for the state. If we decide to add more input steps, we can modify the state machine and add more input steps. We are not required to hold state in the app (it is all delegated to Burr’s persistence), so we can easily load up from any given point, allowing the user to wait for seconds, minutes, hours, or even days before continuing.

As the frontend simply renders based on the current state and the next step, it will always be correct, and the user can always pick up where they left off. With Burr’s telemetry capabilities you can debug any state-related issues, ensuring a smooth user experience.

Adding a UI

Now that we have a set of endpoints, the UI is simple. In fact, it mirrors the API almost exactly. We won’t dig into this too much, but the high-level is that you’ll want the following capabilities:

Render the current state (show the history, latest draft)
Include a form for the next action’s inputs (provide feedback, answer clarifications)
Post the results to your FastAPI endpoints, pause for response, GOTO (1)

You can see the UI here. Here’s an example of it in action:

You can play around with it if you download burr (pip install “burr[start]” && burr), and navigate to http://localhost:7241/demos/email-assistant.

Note that there are many tools that make this easier/simpler to prototype, including chainlit, streamlit, etc… The backend API we built is amenable to interacting with them as well.

Additional Capabilities

Customizing Persistence

While we used the simple SQLLite persister, you can use any of the others that come with Burr or implement your own to match your schema/db infrastructure. To do this you implement the BaseStatePersister class, and add it in with the ApplicationBuilder, instead of the SQLLite persister we used above.

Additional Monitoring/Visibility

Using the Burr UI to monitor is not the only way. You can integrate your own by leveraging lifecycle hooks, enabling you to log data in a custom format to, say, datadog, langsmith, or langfuse.

Furthermore, you can leverage additional monitoring capabilities to track spans/traces, either logging them directly to the Burr UI or to any of the above providers. See the list of available hooks here.

Async/Streaming

While we kept the APIs we exposed synchronous for simplicity, Burr supports asynchronous execution as well. Burr also supports streaming responses for those who want to provide a more interactive UI/reduce time to first token.

So how does it do in practice?

As with any LLM application, the entire prompt matters. If you can provide the right guidance, the results are going to be better than if you don’t. Much like if you are going to instruct a human, more guidance is always better. That said, if you find yourself always correcting some aspect, then changing the base prompt is likely the best course of action. For example, using a single-shot or few-shot approach might be a good choice to try to help instruct the LLM as to what you’d like to see given your specific context.

Post Summary

In this post we discussed how to address some of the challenges around building human-in-the-loop agentic workflows. We ran through an example of making an email assistant using Burr to build and run it as a state machine, and FastAPI to run Burr in a web service. We finally showed how you can extend the tooling we used here for a variety of common production needs — e.g. monitoring & storage.

Additional Resources

Join Burr’s Discord for help or if you have questions!
Burr’s Github repository (if you like it we’d love a ⭐️)
FastAPI guide
Technical deep-dive of building with a web-server on github
Code for the email assistant

Using IPython Jupyter Magic commands to improve the notebook experience

Stefan Krawczyk — Sun, 03 Mar 2024 04:33:55 +0000

Introduction

Jupyter Notebooks are commonplace in data science. They allow a mixture of “read, evaluate, print, [in a] loop” (REPL) code writing and documentation in a single place. They’re most commonly used for analysis and brainstorming purposes, but also, more contentiously, some prefer notebooks to scripts to run production code (but we won’t focus on that here).

Invariably, the code written in notebooks will be repetitive in some way such as setting up a database connection, displaying an output, saving results, interacting with an internal platform tool, etc. It’s best to store this code as functions and/or modules to make them reusable and more easily manage them.

However, the notebook experience is not always improved when you do that. For instance, you still need to import and call these functions throughout your notebook, which doesn’t change the notebook experience much at all. So what is the answer to augmenting the notebook development experience itself? IPython Jupyter Magic Commands.

IPython Jupyter Magic commands (e.g. lines in notebook cells starting with % or %%) can decorate a notebook cell, or line, to modify its behavior. Many are available by default, including %timeit to measure the execution time of the cell and %bash to execute shell commands, and others are provided by extensions such as %sql to write SQL queries directly in a cell of your notebook.

In this post, we’ll show how your team can turn any utility function(s) into reusable IPython Jupyter magics for a better notebook experience. As an example, we’ll use Hamilton, my open source library, to motivate the creation of a magic that facilitates better development ergonomics for using it. You needn’t know what Hamilton is to understand this post.

Note. Nowadays, there are many flavors of notebooks (Jupyter, VSCode, Databricks, etc.), but they’re all built on top of IPython. Therefore, the Magics developed should be reusable across environments.

Understanding IPython Jupyter Magics

IPython Jupyter Magics (which we’ll shorten to just Magics) are bits of code that can be dynamically loaded into your notebooks. They come in two flavors, line and cell magics.

Line magic, as it suggests, operates on a single line. That is, it only takes in as input what is specified on the same line. They are denoted by a single % in front of the command.

# will only time the first line
%time print(“hello”)
print(“world)

Cell magic, as it suggests, takes the entire contents of a cell. They are denoted by a double %% in front of the command.

# will time the entire cell
%%timeit
print(“hello”)
print(“world”)

Jupyter comes with several in-built magic commands. You can think of them as “command line” tools that have access to the entire notebook context. This allows them to interact with the notebook output (e.g., printing results, displaying a PNG, rendering HTML), but also to modify the state of existing variables and write to other code and markdown cells!

This is great for developing internal tooling because it can abstract away and hide from the user unnecessary complexities, making the experience “magical”. This is a powerful tool to develop your own “platform efforts”, especially for MLOps and LLMOps purposes, as you can hide what is being integrated with from having to be exposed in the notebook. It therefore also means that notebooks don’t need to be updated if this abstracted code changes under-the-hood, since it can all be hidden in a python dependency upgrade.

Workflow example

Magic commands have the potential to make your workflow simpler and faster. For example, if you prefer to develop in a notebook before moving your code to a Python module, this can involve error-prone cutting & pasting. For this purpose, the magic %%writefile my_module.py will directly create a file and copy your cell content into it.

On the opposite hand, you might prefer developing in my_module.py in your IDE and then load it into a notebook to test your functions. This usually involves restarting the notebook kernel to refresh the module imports, which can be tedious. In that case, %autoreload will automatically reload every module imports before every cell execution, removing this friction point!

The Need for Custom Magic Commands

In the post How well-structured should your data code be?, it is argued that standardization/centralization/“platform” efforts, should change the shape of the “move quickly vs. built-to-last” trade-off curve for the better. A concrete tactic to change this trade-off is to implement better tooling. Better tooling should make what used to be complex, simpler and accessible. Which is exactly what you can achieve with your own custom Magic commands, which translates to less of a trade-off to be made.

A Hamilton Magic Command

For those unfamiliar with Hamilton we point readers to the many articles on it (e.g. origin story, ML pipeline reference, Production prompt engineering, etc.) as well as https://www.tryhamilton.dev/.

Hamilton is an open source tool that we created at Stitch Fix back in 2019. Hamilton helps data scientists and engineers define testable, modular, self-documenting dataflows, that encode lineage and metadata. Hamilton achieves these traits in part by requiring Python functions to be curated into modules.

Nonetheless, the typical Jupyter notebook usage pattern leads to code residing in the notebook and nowhere else posing a developer ergonomics challenge:

How can we enable someone to create Python modules easily and quickly from a notebook, while also improving the development experience?

The Hamilton developer loop looks like the following:

Hamilton development loop. Image by author.

Take a minute to read this loop. The loop shows that anytime a code change is made, the user would need to not only re-import the Python module, but also recreate the Driver object as well. Since notebooks allow cell execution in any order, it can become difficult for the user to track which version is loaded for each module and what is currently loaded in a Driver. This burden lies on the user and might require restarting the kernel, which would lose other computations (thankfully, Hamilton can be set up to execute complex dataflows and resume where you left-off…), which is less than ideal.

Here’s how we could improve this loop using Magics:

Create a “temporary” Python module from the functions defined in a cell, and import this new module directly in the notebook.

Automatically visualize the directed acyclic graph (DAG) defined by the functions to reduce visualization boilerplate code.

Rebuild all Hamilton Drivers found in the notebook with updated modules, saving the user time to have to remember to manually recreate drivers to pick up the change.

What we’re going to implement

We would like a command that looks like this:

%%cell_to_module -m my_module --display --rebuild-drivers

def my_func(some_input: str) -> str:
  """Some logic"""
  return ...

And cause the following behavior after running the cell:

Create a module with the name my_module in the notebook.
Display the DAG constructed by the functions within the cell.
Rebuild any downstream drivers that used my_module in other cells, saving the user having to re-run those cells.

As you can see this is a non-trivial Magic command, since we’re adjusting output of the cell and state of the notebook.

Building a Custom Magic Command

Here we lay out step by step how to create a Magic Command. To avoid only showing a trivial “hello world” example, we’ll explain how we built Hamilton’s %%cell_to_module magic as well.

Step 1: Setting up your code

Create a new Python module where we’ll write the magic code and a jupyter notebook to try it. The name of this module (i.e., .py file) will be the name of the extension you’ll need to load.

If Jupyter notebook is installed, you have all the required Python dependencies already. Then, add libraries you will need for your Magic, in our case Hamilton (pip install sf-hamilton[visualization]).

Step 2: Define your Magic Command

To define a simple Magic command, you can use functions or objects (see these docs). For more complex Magics where state is needed, you will need the class approach. We’ll use the class based approach here. To start we need to import IPython modules/functions and then define a class that inherits magic.Magics. Each method decorated with @cell_magic or @line_magic defines a new magic, and the class can house however many of them.

To start, your code should look like this at a high level:

# my_magic.py

from IPython.core import magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring

@magic.magics_class
class MyMagic(magic.Magics):
   """Custom class you write"""

   @magic_arguments() # needs to be on top to enable parsing
   @argument(...)   
   @magic.cell_magic
   def a_cell_magic_command(self, line, cell):
      ...

   @magic_arguments() # needs to be on top to enable parsing
   @argument(...)   
   @magic.line_magic
   def a_line_magic_command(self, line):
     ...

For stateful magic, it can be useful to add an __init__() method (i.e. constructor). It is not needed in our case.

By inheriting from magic.Magics, this class has access to several important fields including self.shell, which is the IPython InteractiveShell that underlies the notebook. Using it allows you to pull and introspect variables loaded in the active Jupyter notebook.

Our Hamilton Magic Command will start off looking like:

from IPython.core import magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring

@magic.magics_class
class HamiltonMagics(magic.Magics):
    """Magics to facilitate Hamilton development in Jupyter notebooks"""

  @magic_arguments()  # needed on top to enable parsing
  @arguments(...)
  @magics.cell_magic
  def cell_to_module(self, line, cell):
    ...

Step 3: Parsing input arguments

Next, we specify what arguments will be passed and how to parse them. For each argument, add a @argument, and add a @magic_arguments() decorator on top. They follow a similar pattern to argparse arguments if you’re familiar, but they aren’t quite as fully featured. Within the function, you need to call the parse_argstring() function. It receives the function itself to read instructions from the decorators, and line (the one with % or %%) which contains the arguments values.

Our command would start to look like this:

   @magic_arguments() # needs to be on top to enable parsing
   # flag, long form, default value, help string.
   @argument("-a", "--argument", default="some_value", help="Some optional line argument")   
   @magic.cell_magic
   def a_cell_magic_command(self, line, cell):
       args = parse_argstring(self.a_cell_magic_command, line)

       if args.argument:
          # do stuff -- place your utility functions here

Note, for required arguments, there is no facility in magic_arguments() for that, so you need to manually check in the body of the function for correctness, etc.

Continuing our dissection of the Hamilton Magic example, the method on the class now looks like the following; we use many optional arguments:

    @magic_arguments()  # needed on top to enable parsing
    @argument(
        "-m", "--module_name", help="Module name to provide. Default is jupyter_module."
    )  # keyword / optional arg
    @argument(
        "-c", "--config", help="JSON config, or variable name containing config to use."
    )  # keyword / optional arg
    @argument(
        "-r", "--rebuild-drivers", action="store_true", help="Flag to rebuild drivers"
    )  # Flag / optional arg
    @argument(
        "-d", "--display", action="store_true", help="Flag to visualize dataflow."
    )  # Flag / optional arg
    @argument(
        "-v", "--verbosity", type=int, default=1, help="0 to hide. 1 is normal, default"
    )  # keyword / optional arg
    @magics.cell_magic
    def cell_to_module(self, line, cell):
        """Execute the cell and dynamically create a Python module from its content.

        A Hamilton Driver is automatically instantiated with that module for variable `{MODULE_NAME}_dr`.

        > %%cell_to_module -m MODULE_NAME --display --rebuild-drivers

   Type in ?%%cell_to_module to see the arugments to this magic.

        """
        # specify how to parse by passing 
        args = parse_argstring(self.cell_to_module, line)  
        # now use args for logic

Note, the extra arguments to @argument are useful for when someone uses ? to query what the magic does. I.e. ?%%cell_to_module will show documentation.

Step 4: Implement the Command’s Logic

Now that we have parsed the arguments, we can implement the logic of the magic command. There aren’t any particular constraints here and you can write any Python code. Skipping a generic example (you have enough to get started from the prior step), let’s dig into our Hamilton Magic example. For it, we want to use the arguments to determine the desired behavior for the command:

Create the Python module with module_name.
If —-rebuild-driver, rebuild the drivers, passing in verbosity.
If —-config is present, get that ready.
If —-display, display the DAG.

See comments in the code for explanations:

# we're in the bowels of def cell_to_module(self, line, cell):
# and we remove an indentation for readability
...
# specify how to parse by passing this method to the function 
args = parse_argstring(self.cell_to_module, line)

# we set a default value, else use the passed in value 
# for the module name.
if args.module_name is None:
    module_name = "jupyter_module"
else:
    module_name = args.module_name

# we determine whether the configuration is a variable
# in the notebook environment
# or if it's a JSON string that needs to be parsed.
display_config = {}
if args.config:
    if args.config in self.shell.user_ns:
        display_config = self.shell.user_ns[args.config]
else:
    if args.config.startswith("'") or args.config.startswith('"'):
        # strip quotes if present
        args.config = args.config[1:-1]
    try:
        display_config = json.loads(args.config)
    except json.JSONDecodeError:
        print("Failed to parse config as JSON. "
              "Please ensure it's a valid JSON string:")
        print(args.config)

# we create the python module (using a custom function)
module_object = create_module(cell, module_name)

# shell.push() assign a variable in the notebook. 
# The dictionary keys are the variable name
self.shell.push({module_name: module_object})

# Note: self.shell.user_ns is a dict of all variables in the notebook
# -- we pass that down via self.shell.
if args.rebuild_drivers:
    # rebuild drivers that use this module (custom function)
    rebuilt_drivers = rebuild_drivers(
        self.shell, module_name, module_object, 
        verbosity=args.verbosity
    )
    self.shell.user_ns.update(rebuilt_drivers)

# create a driver to display things for every cell with %%cell_to_module
dr = (
    driver.Builder()
    .with_modules(module_object)
    .with_config(display_config)
    .build()
)
self.shell.push({f"{module_name}_dr": dr})

if args.display:
    # return will go to the output cell. 
    # To display multiple elements, use IPython.display.display(
    # print("hello"), dr.display_all_functions(), ... )
    return dr.display_all_functions()

Notice how we use self.shell. This allows us to update and inject variables into the notebook. The values returned by the function will be used as “cell output” (where you see printed values).

Step 5: Register your Magic Command

Lastly, we need to tell IPython and the notebook about the Magic Command. Our module where our Magic is defined must have the following function to register our Magic class, and be able to load our extension. If doing anything stateful, this is where you’d instantiate it.

Notice that the argument ipython here is the same InteractiveShell available through self.shell in the class method we defined.

def load_ipython_extension(ipython: InteractiveShell):
  """
  Any module file that define a function named `load_ipython_extension`
  can be loaded via `%load_ext module.path` or be configured to be
  autoloaded by IPython at startup time.
  """
  ipython.register_magics(MyMagic)
  ipython.register_magics(HamiltonMagics)

See the full Hamilton Magic Command here.

Try it out

To load your magic in the notebook, try the following:

%load_ext my_magic

in the case of our Hamilton Magic we'd load it via:

%load_ext hamilton.plugins.jupyter_magic

While you’re developing, use this to reload your updated magic without having to restart the notebook kernel.

%reload_ext my_magic

You can then invoke the magic commands defined on a per line or cell basis. So for the Hamilton one we’d now be able to do:

?%%cell_to_module

Here’s an example use of it, with it injecting the visualization:

Animated gif of adding functions and hitting enter to refresh the image.

What you’d do next

In a real world use case, you would most likely version and package your magic likely into a library, that you can then manage easily in python environments as required. With the Hamilton Magic Command it’s packaged into the Hamilton library, and thus to get it, one need only install sf-hamilton and loading the magic command would become accessible in the notebook.

To close

In this post we showed you the steps required to create and load your own IPython Jupyter Magic Command. Hopefully you’re now thinking of the common cells/task/actions that you perform in a notebook setting, which could be enhanced/simplified/or even removed with the addition of a simple Magic!

To demonstrate a real-life example, we motivated and showed the internals of a Hamilton Magic Command to show a command that was built to improve the developer experience within a Jupyter notebook, by augmenting the output and changing internal state.

We hope that this post helps you get over the hump and create something more ergonomic and useful for you and your teams’ Jupyter Notebook experience.

DEV Community: Stefan Krawczyk

Building an Email Assistant Application with Burr

Why are interactive agents applications a challenge?

The Tools

Burr

FastAPI

React (or any frontend framework)

Building

Modeling Control Flow

Implementing/Testing

Persistence

Integrating in a web server

Endpoints

Endpoints

Adding a UI

Additional Capabilities

Customizing Persistence

Additional Monitoring/Visibility

Async/Streaming

So how does it do in practice?

Post Summary

Additional Resources

Using IPython Jupyter Magic commands to improve the notebook experience

Introduction

Understanding IPython Jupyter Magics

Workflow example

The Need for Custom Magic Commands

A Hamilton Magic Command

What we’re going to implement

Building a Custom Magic Command

Step 1: Setting up your code

Step 2: Define your Magic Command

Step 3: Parsing input arguments

Step 4: Implement the Command’s Logic

Step 5: Register your Magic Command

Try it out

What you’d do next

To close

Some other resources