Composable AI systems: Building and AI Stylist Specialist selling our productsWhat's a composable AI system
LLMs are great at answering general questions. However, general intelligence alone isn't enough to provide value to your customers.
To be able to provide valuable answers, extra information is required, specific to your business and the user asking the question (your customer contract ID, the last email they sent to your support, your most recent sales report etc.).
Composable AI systems are designed to answer this challenge. They are more advanced AI deployments, composed of multiple entities (tools) specialized in different action (retrieving information or acting on external systems).
At a high level, you build & present a set of custom functions to the AI. The LLM can then reason about it, deciding which tool should be called and information gathered to answer the customer need.
Building Composable AI Systems with Databricks Mosaic AI agent framework
Databricks simplifies this by providing a built-in service to:
- Create and store your functions (tools) leveraging UC
- Execute the functions in a safe way
- Reason about the tools you selected and chain them together to properly answer your question.
At a high level, here is the AI system we will implement in this demo:
Creating our tools: Using Unity Catalog Functions
Let's start by defining the functions our LLM will be able to execute. These functions can contain any logic, from simple SQL to advanced python.
Computing Mathematics: converting inches to centimeters
Our online shop is selling across all regions, and we know our customers often ask to convert cm to inches. However, plain LLMs are quite bad at executing math. To solve this, let's create a simple function doing the conversion.
We'll save this function within Unity Catalog. You can open the explorer to review the functions created in this notebook.
*Note: This is a very simple first example for this demo. We'll implement a broader math tool later on.*
Executing a tool to fetch internal data: getting the latest customer orders
We want our stylist assistant to be able to list all existing customer orders
Executing a Python function to fetch external dataset in realtime: getting the weather
We want our stylist assistant to give us recommendations based on the weather. Let's add a tool to fetch the weather based on longitude/latitude, using Python to call an external Weather API.
**Note: This will be run by a serverless compute, and accessing external data, therefore requires serverless network egress access. If this fails in a serverless setup/playground, make sure you are allowing it in your networking configuration (open the workspace network option at admin account level)**
Creating a function calling LLMs with specific prompt as a tool
You can also register tools containing custom prompts that your LLM can use to to execute actions based on the customer context.
Let's create a tool that recommend the style for our user, based on the current weather.
Using Vector search to find similar content as a tool
Let's now add a tool to recomment similar items, based on an article description.
We'll be using Databricks Vector Search to perform a realtime similarity search and return articles that we could suggest from our database.
To do so, you can leverage the new `vector_search` SQL function. See [Documentation](https://docs.databricks.com/en/sql/language-manual/functions/vector_search.html) for more details.
To simplify this demo, we'll fake the call to the vector_search and make it work without pre-loading the demo
Using Databricks Playground to test our functions
Databricks Playground provides a built-in integration with your functions. It'll analyze which functions are available, and call them to properly answer your question.
To try out our functions with playground:
- Open the [Playground](/ml/playground)
- Select a model supporting tools (like Llama3.1)
- Add the functions you want your model to leverage (`catalog.schema.function_name`)
- Ask a question (for example to convert inch to cm), and playground will do the magic for you!
Note: Tools in playground is in preview, reach-out your Databricks Account team for more details and to enable it.
Building an AI system leveraging our Databricks UC functions with Langchain
These tools can also directly be leveraged on custom model. In this case, you'll be in charge of chaining and calling the functions yourself (the playground does it for you!)
Langchain makes it easy for you. You can create your own custom AI System using a Langchain model and a list of existing tools (in our case, the tools will be the functions we just created)
Enable MLflow Tracing
Enabling MLflow Tracing is required to:
- View the chain's trace visualization in this notebook
- Capture the chain's trace in production via Inference Tables
- Evaluate the chain via the Mosaic AI Evaluation Suite
Start by creating our tools from Unity Catalog
Let's use UCFunctionToolkit to select which functions we want to use as tool for our demo:
Let's create our langchain agent using the tools we just created
Let's give it a try: Asking to run a simple size conversion.
Under the hood, we want this to call our Math conversion function:
Chaining multiple tools
Let's see if the agent can properly chain the calls, using multiple tools and passing the values from one to another. We'll ask for an outfit for today, and the expectation is that the LLM will fetch the weather before.
Extra: running more advanced functions using python:
Let's see a few more advanced tools example
Calculator tool: Supporting math operations
Let's add a function to allow our LLM to execute any Math operation.
Databricks runs the python in a safe container. However, we'll filter what the function can do to avoid any potential issues with prompt injection (so that the user cannot execute other python instructions).
Run any python function
If we are creating a coding assistant, it can be useful to allow our AI System to run python code to try them and output the results to the user. **Be careful doing that, as any code could be executed by the user.**
Here is an example:
Deploying our Agent Executor as Model Serving Endpoint
We're now ready to package our chain within MLFLow, and deploy it as a Model Serving Endpoint, leveraging Databricks Agent Framework and its review application.
Deploying the agent as a model serving endpoint
Databricks automatically generates all the required notebooks and setup for you to deploy these agents!
To deploy them with the latest configuration, open [Databricks Playground ](/ml/playground), select the tools you want to use and click on the "Export Notebook" button on the top!
This will generate notebooks pre-configured for you, including the review application to start testing and collecting your new model!
Conclusion
That's it, we saw how you can leverage Databricks Functions as tools within your AI system.
This demo gave you a simple overview and some examples to get started. You can take it to the next level within your own AI Systems!
Once your application is complete and your chain ready to be deployed, you can easily serve your model as a Model Serving Endpoint, acting as your Compound AI System!
Coming soon
We'll soon update this demo to add more examples, including:
- how to properly leverage secrets within AI functions to call external systems having Authentication
- generating images and calling ML model endpoints
- Use Genie as an agent to query any table!
- More to come!
Stay tuned!