AI-powered applications are all the buzz now and we too want to surf this wave of hype. What better way is there than to build an AI app from top to bottom in Scala - a full product, not just the app itself but the whole thing - including an infrastructural blueprint necessary to deploy it to the live environment.
What would be a suitable objective then? Something that would be useful and interesting - a self-hosted AI-based open source documentation assistant maybe? If we could host it on the cheap somewhere it would also be nice.
The architecture of an AI app in Scala 3 and Besom
For the AI app itself we're going to use tapir and sttp (with an openAI extension) for http, magnum for database access and Loom to make things easier.
Why do we need database access you might ask? Ah, that's because we're going to do machine learning the best way possible - via Postgres! PostgresML is a Postgres distribution that brings ML functionality to the SQL level where we can easily work with it via JDBC without dealing with the whole Python / C mess underneath.
The AI app will utilize Retrieval Augmented Generation pattern and leverage PostgresML built-in vector database functionality to store information from our docs as embeddings. This will allow us to find things our users ask about in our db using cosine similarity and then use the best semantically matching text as part of our crafted prompt to GPT 3.5 Turbo.
Since we're going to run containers and since we want to emulate the best practices of the startup world (mostly the overengineering part!) we're going to deploy a Kubernetes cluster. To avoid running over our non-existent budget, we'll avoid pricy hyperscalers and go with the cheapest provider that also works nicely with Pulumi - Hetzner.
Hetzner's offering does not offer a managed Kubernetes, but that shouldn't stop us! We could utilise the awesome power of Besom and Pulumi to just deploy a full-blown Kubernetes distribution like Typhoon. But we have to think like a startup and cut corners a bit so we'll just go with Rancher's K3S.
Here's a chart of how the infrastructure of the whole system will look like - a 3-node k3s cluster with a built-in Traefik ingress and all the necessary Kubernetes entities.
Let’s build the AI App in Scala 3
Our AI app itself will be quite simple - a plain http server serving pure html from index endpoint with a single POST /inquire endpoint to handle queries from the user. To be able to answer our users' questions about Besom we will have to be able to find relevant docs and to do that we need to create and populate embeddings table in PostgresML.
As this is a PoC (we're doing a startup, remember?) we will just put our markdown files into a directory that we will include on the classpath using scala-cli's directive: //> using resourceDir content. Then a small migrations script will create tables and populate them. Interesting bit here is how easy it is to create embeddings using PostgresML:
This single query uses data stored in docs table to create a parallel table enriched with embeddings using PostgresML's built-in function pgml.embed with a Microsoft's E5 embedding model. The query side looks like this:
This time we use a common table expression (WITH) to create an embedding of the query from our user and then in the main query we compare the resulting vector with embeddings stored in our docs_embeddings table. After ordering by result of that comparison in descending order we are able to fetch the doc that matched the query the most.
The next step is to pass the doc along with it's url and the original query of the user to a large language model that will compose a response. We do this using the sttp-openai module to call GPT 3.5 Turbo with our custom prompt via OpenAI API. GPT 3.5 Turbo is both cheap and relatively fast so should be more than enough for our purposes for now:
The last thing to do is to convert AI model's response to HTML as it will most surely respond in markdown (we don't even have to specify that in our prompt). We will use Flexmark library to do this:
Finally http layer with tapir is just two endpoints and a simple, plain http server that works well with Loom's virtual threads:
We can skip the Templates object as it only hosts two functions that return Strings of HTML. We also expose an endpoint to server static resources from the classpath as we have a nifty, animated loading indicator to show to our users while AI is crafting responses to their queries.
Publishing the AI app container
To get our AI app ready for deployment we have to publish it as a docker container available for download from a public docker container registry. Fortunately for us Github offers a container registry for its users. We have to start with logging into ghcr.io with our local docker to be able to push there:
We can use Github's gh command line tool to fetch a token for us to streamline the process or just create a general token with write permissions to Github Container Registry in the web UI.
The next step is to build our AI app with scala-cli:
Then we have to build an image so let's start with a Dockerfile:
and then let's use it to build the image and push it to ghcr.io:
Should you want to follow the same steps, remember to use your own Github username here.
Having built the AI app we still need to run it somewhere and that place has also allow us to deploy a PostgresML instance. Let's create our own platform!
The AI app’s infrastructure with Besom
First things first - we need servers. To create them on Hetzner Cloud we need to provide our public ssh key and few additional configuration parameters to create three servers in three different locations:
We load the path to our public and private (we will need it later) ssh keys from the Pulumi configuration and then create a SshKey resource in Hetzner. Afterwards we create 3 servers running Ubuntu 22.04 of cx21 type with 2 vCPUs and 4 GB of RAM.
Finally we make sure they get created in parallel by turning Vector[Output[hcloud.Server]] into Output[Vector[hcloud.Server]] using the parSequence extension and putting resulting Output into Stack's dependencies. We also export ip addresses of all three nodes just to be able to see them as results of our pulumi up.
With this done we only need to provide the following three configuration properties.
- ssh_public_key_path
- ssh_private_key_path
- hcloud_token
using pulumi config set while remembering to use --secret for all secrets and then we can apply the stack and connect to any of our nodes via ssh because we've already set them up with our public ssh key.
Deploying Kubernetes for our AI app
The next step is to deploy K3S onto our servers to get a unified orchestration engine over all our machines. To do this we will leverage the Pulumi command package that allows us to run arbitrary commands on both local and remote machines.
The procedure itself is quite simple, but we have to translate each step where we run remote commands to Pulumi's CRUD-like resource model. The actions we have to complete are:
- Initialize a cluster leader node that will provide us with a join token that we have to use on follower node startup to form a cluster
- Initialize follower nodes with the aforementioned join token
- configure all the nodes with our Github Container Registry token that will allow us to pull images of our AI app published to ghcr.io
- Fetch local kubeconfig file from the leader node and apply some small tweaks to it so that we can use it to execute kubectl commands remotely
The code to do that will look like this:
The interesting bit here is that given our commands are often unrelated on data level (outputs of previous are not inputs of the next) so we have to establish order of execution manually using the dependsOn resource option.
Making of a Component
Moreover, this bit of code looks like it's reusable and we have already defined a K3S case class to keep interesting outputs of this block so maybe it's worth turning it into a reusable component. To do that we need to provide some boilerplate data structures that form the API we know from other provider packages:
This is the first time we need to work with Inputs to define a funnel-like API that allows others to pass values that are raw literals like String or dynamic values pulled from configuration or other resources as, for example, Output[String]. Inputs are in fact opaque types over union types and the most basic Input[A] is equivalent to A | Output[A] allows one to pass either a value of A type or its dynamic variant wrapped in Output. The family of extension functions seen here as .asOutput() or .asManyOutput() does the safe conversion from the opaque type into a lifted form, e.g.: Input[String] (so, as we now know String | Output[String]) into Output[String].
Taking care of the plans
One interesting outlier here is that we take serverIps argument as Vector[Input[String]] which yields a Vector[Output[String]] in K3SArgs case class. Why don't we just take an Input[List[String]] or Input.OneOrList[String] instead? Well, we know beforehand that there are just two ways one can obtain a list of ip addresses - one is that the user provides it by hand and a Vector[Input[String]] will accept a Vector[String] also so we're not making it troublesome really. The other way is to take those values from properties of other resources like we're actually doing here by taking ipv4Address of hcloud.Server resource. We have to remember however that properties of resources are unknown in dry runs so if we want to avoid spoiling deployment plan shown in pulumi up we have to avoid propagating unknown status and that's exactly what happens should we use an Output resulting from this computation: serverPool.map(_.ipv4Address).parSequence.
The final reformulation of this code into a Besom component looks like this:
Finally we can use our new component in the main Pulumi.run block:
This in turn looks pretty straightforward and allows us to hide a lot of internal complexity of the deployment beyond a familiar facade of a resource. We also use the deletedWith resource option here to let Pulumi know that removal of this whole layer can be skipped should the leader node be deleted too because that would mean we're destroying the whole thing and there's no point in cleaning up Kubernetes deployment if the underlying server is also getting axed.
Deploying the AI app
There are few things we have to deploy to Kubernetes for our AI app to work: a PostgresML instance, our app itself and some networking cruft like Service and Ingress to be able to route external traffic to it. Having already grokked components why don't we try and make it an encapsulated, reusable piece of infrastructure too? Let's start with the boilerplate of configuration classes:
We also need a beefy import block for all the stuff were going to use:
Now that we have this out of our way we can declare our component:
and start filling it with deployment logic. We are going to need some common constants, create a Kubernetes namespace and extract the values out of our component arguments:
First important thing to do is to deploy the PostgresML instance and to do it nicely with persistent volumes taken care for us is to use a StatefulSet:
We configure our database container with proper readiness and liveness probes, expose ports, set up a custom timeout for creation to ten minutes because the image for PostgresML is huge and deal with the fact that without overridden args the container quits immediately with a good, old tail -f /dev/null. With StatefulSet done we have to create a Service to expose the database to our pods:
We pass all the necessary properties and mark the Service dependent on the creation of the StatefulSet via dependsOn resource option as they are unrelated entities otherwise and Service could fail with a timeout error waiting on the pods of StatefulSet to become available. We then extract the name of the service to get an internal Kubernetes domain name via which we can reach the database and create the first output value of our component.
Next step is to declare a Deployment of our AI app.
There's nothing really fancy happening here - it's a plain, usual specification of a Deployment with configuration of the AI app passed via environment variables, with readiness and liveness probe for it to be a good neighbour in k8s environment and exposed ports. Our Deployment will need a companion Service and an Ingress based on K3S built-in Traefik:
With this done the only thing to do is to compute the second output property of our AppDeployment component:
It's time to use the component in our main Pulumi.run block. We also need to create a Kubernetes Provider instance first to let Pulumi know onto which cluster we want to deploy all of our resources:
The domain being used here is privately owned and is managed by Cloudflare and this fact will come handy in a moment. We set up some configuration arguments and also inform Pulumi that, again, it's not worth dealing with deletion of Kubernetes entities if the whole Kubernetes cluster is going down.
We also do something quite significant here: we pass the k3sProvider as one of ComponentResourceOptions providers. You might have noticed above that we didn't pass opts (provider = k3sProvider) to all of the Kubernetes resources we created - that's because components have a very useful property - all resources belonging to the component inherit it's providers if they are any that match their package name.
This means that all the Kubernetes resources declared in the AppDeployment component will automatically use the correct provider, which is nice and definitely helps to keep bugs away.
Going public
Final step of the whole thing is to tinker a bit with the domain controlled by Cloudflare to make it point towards our cluster's and it's ingress ports. This bit is small enough not to put it into a component so we're going to just add it to the main Pulumi.run block as is:
We create a Provider instance for Cloudflare and then use it to create an A DNS record for each of the servers. We also configure proxying setup so that we can benefit from Cloudflare's automatic TLS setup (provided it's set to flexible mode!) without deploying our own CertManager to set up Let's Encrypt certs. That's obviously something we should do in the future but for now this is good enough.
There's also a small helper piece that will help us connect to the newly deployed cluster which basically writes the adjusted kubeconfig.conf file into our working directory:
We modify the final Stack definition:
and our full deployment is ready. Now we can trigger pulumi up and see that everything deploys just fine beside our AI app, which ends up in a crash loop. We fetch the logs from the pod and see that we forgot to pass the PORT env variable. Sigh, such a dumb error, if only we used a language with superpowers...
Using Scala 3 superpowers for safety and profit
The piece of code that caused a small headache for us is this part of the deployment of our AI app.
One of the EnvVarArgs is missing and it's very easy to miss this because there's an impassable border between our AI app, which is currently nicely wrapped in a docker image, and the runtime of infrastructure code. Is it really impassable? It turns out that if you use Scala there aren't many things that are impossible to do! Enter the besom-cfg package which allows us to safely connect the configuration of applications with the declarations of the infrastructure layer that will provide necessary configuration. Let's see how this works and start with an assumption that we're using a naive, hand-written configuration resolution:
It's worth to note that this could be done with pureconfig, ciris or zio-config - it doesn't matter as these libraries do exactly the same thing: resolve configuration from configuration files and/or environment variables without exposing any information about what configuration variables are necessary for the AI app to start successfully.
That's not the case with besom-cfg as the resolution of configuration case classes known from aforementioned libraries is just a part of what it does. More important bit is that it allows the infrastructural Besom program to introspect into the AI app at compile time and exposes information about what is really necessary for the app to successfully start. Let's see how this looks like in the application code:
That's all, just derives Configured clause on a case class with our configuration and a call to resolveConfiguration function to load it. Nothing else has to be done. Now let's see what has to change in infrastructural code. We need to include some new imports:
and then replace things in the deployment of our AI app.
Notice that we're still missing port property that is a part of the Config class on the application's side but now this is a compile error of the infrastructure program and this successfully prevents such a mistake:
besom-cfg performs type checking of the configuration and will be capable of passing information about type refinements from libraries like Iron too:
The library currently supports Kubernetes containers only but other forms of deployment and integrations will be available in the future.
Deployment time!
After the setup of besom-cfg and correction of configuration for the AI app it's finally time to deploy the whole stack. To do this we have to provide the rest of configuration properties using pulumi config set:
- cloudflare_token
- cloudflare_zone_id
- github_docker_token
- openai_token
After this is done we can execute the deployment:
Now we can visit the deployed site at the provided url and ask some questions about Besom.
Conclusion
The AI app itself is obviously a very limited PoC built to impress readers and it would require a lot of further iterations to make it really useful (along with a widget to embed it in an actual docs site) but it proves the point - full stack Scala is really here. It also hopefully conveys that it’s possible to build experiments, proofs of concepts and small apps in Scala 3 with code that’s easy to write, easy to grasp and that still offers much more scalability and type safety than alternatives. Moreover, with new developments around Besom, Scala offers a productive environment that integrates application layer with infrastructure layer and removes whole classes of errors in the deployment phase of the application lifecycle.
Full code for this blogpost is available here: https://github.com/VirtusLab/besom-ask-me
Our deployment will be available for some time but feel free to deploy your own!