Rust with axum

Cover Page

Back-end Page

We assume that your chatterd code base has accumulated code up to at least the llmChat back end.

toolbox

Let us start by creating a toolbox to hold our tools. Change to your chatterd folder and create a new Rust file, name it toolbox.rs:

server$ cd ~/reactive/chatterd 
server$ vi src/toolbox.rs

Put the following use imports at the top of the file:

#![allow(non_snake_case)]
use crate::AppState;
use futures::future::{BoxFuture, FutureExt};
use serde::{Deserialize, Serialize};
use serde_json::from_str;
use std::collections::HashMap;
use std::sync::LazyLock;
use std::borrow::Cow;

The contents of this file can be categorized into three purposes: tool/function definition, the toolbox itself, and tool use (or function calling).

Tool/function definition

Ollama tool schema: at the top of Ollama’s JSON tool definition is a JSON Object respresenting a tool schema. The tool schema is defined using nested JSON Objects and JSON Arrays. Add the full nested definitions of Ollama’s tool schema to your file:

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OllamaToolSchema {
    #[serde(rename = "type")]
    type_: Cow<'static, str>,
    function: OllamaToolFunction,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct OllamaToolFunction {
    name: Cow<'static, str>,
    description: Cow<'static, str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    parameters: Option<OllamaFunctionParams>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct OllamaFunctionParams {
    #[serde(rename = "type")]
    type_: Cow<'static, str>,
    properties: HashMap<String, OllamaParamProp>,   // HashMap has no ordering
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    required: Vec<String>,                          // parameters MUST be in function-signature order
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct OllamaParamProp {
    #[serde(rename = "type")]
    type_: Cow<'static, str>,
    description: Cow<'static, str>,
    #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
    enum_: Option<Vec<Cow<'static, str>>>, // `enum` is keyword in Rust
}

Weather tool schema: in this tutorial, we have only one tool resident on the back end, get_weather. Instead of manually instantiating an OllamaToolSchema for each tool, we use Rust’s serde package to create one for us from a JSON schema file. We’ll use the include_str! macro to read in the file and feed it directly to serde::from_str later when populating the TOOLBOX.

Weather tool function: we implement the get_weather tool as a getWeather() function that makes an API call to the free Open Meteo weather service. Add the following struct definitions to hold Open Meteo’s return result. For this tutorial, we’re only interested in the latitude, longitude, and temperature returned by Open Meteo:

#[derive(Deserialize)]
struct Current {
    #[serde(rename = "temperature_2m")]
    temp: f64,
}

#[derive(Deserialize)]
struct OMeteoResponse {
    latitude: f64,
    longitude: f64,
    current: Current,
}

Here’s the definition of the getWeather() function:

pub async fn getWeather(appState: &AppState, argv: &[String]) -> Result<Option<String>, String> {
    // Open-Meteo API doc: https://open-meteo.com/en/docs#api_documentation
    match appState.client
        .get(format!("https://api.open-meteo.com/v1/forecast?latitude={}&longitude={}&current=temperature_2m&temperature_unit=fahrenheit",
                     argv[0], argv[1]))
        .send().await {
        Ok(response) => {
            let ometeoResponse: OMeteoResponse = response.json().await.unwrap();
            Ok(Some(format!("Weather at lat: {}, lon: {} is {}ºF",
                            ometeoResponse.latitude, ometeoResponse.longitude, ometeoResponse.current.temp)))
        },
        Err(err) => {
            Err(format!("Open-meteo: {}", err))
        }
    }
}

The toolbox

Even though we have only one resident tool in this tutorial, we want a generalized architecture that can hold multiple tools and invoke the right tool dynamically. To that end, we’ve chosen to use a switch table (or jump table or, more fancily, service locator registry) as the data structure for our tool box. We implement the switch table as a dictionary. The “keys” in the dictionary are the names of the tools/functions. Each “value” is a record containing the tool’s definition/schema and a pointer to the function implementing the tool. To send a tool as part of a request to Ollama, we look up its schema in the switch table and copy it to the request. To invoke a tool called by Ollama in its response, we look up the tool’s function in the switch table and invoke the function.

Add the following type for an async tool function and the record type containing a tool definition and the async tool function:

type ToolFunction = 
    dyn for <'a> Fn(&'a AppState, &'a [String]) -> BoxFuture<'a, Result<Option<String>, String>> + Send + Sync;

pub struct Tool {
    pub(crate) schema: OllamaToolSchema,
    function: Box<ToolFunction>,
}

Now create a switch-table toolbox and populate it with the weather tool using the include_str! macro to read get_weather.json schema file and feed it directly to serde::from_str to be deserialized into an instance of OllamaToolSchema (the type of the schema property):

pub const TOOLBOX: LazyLock<HashMap<String, Tool>> = LazyLock::new(|| {
    HashMap::from([
        (
            "get_weather".to_string(),
            Tool {
                schema: from_str(include_str!("../tools/get_weather.json"))
                    .expect("Failed to parse get_weather tool schema"),
                function: Box::new(|appState, argv| {
                    async { getWeather(appState, argv).await }.boxed()
                }),
            },
        ),
    ])
});

Tool use or function calling

Ollama tool call: Ollama’s JSON tool call comprises a JSON Object containing a nested JSON Object carrying the name of the function and the arguments to pass to it. Add these nested struct definitions representing Ollama’s tool-call JSON to your toolbox.rs file:

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OllamaToolCall {
    pub(crate) function: OllamaFunctionCall,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OllamaFunctionCall {
    pub(crate) name: String,
    arguments: HashMap<String, String>,
}

Tool invocation: finally, here’s the tool invocation function. We call this function to execute any tool call we receive from Ollama response. It looks up the toolbox for the tool name. If the tool is resident, it runs it and returns the result, otherwise it returns a null.

pub async fn toolInvoke(
    appState: &AppState,
    function: OllamaFunctionCall,
) -> Result<Option<String>, String> {
    if let Some(tool) = TOOLBOX.get(&function.name) {
        // get arguments in order, they may arrive out of order from Ollama
        let argv: Vec<String> = tool.schema.function.parameters
            .as_ref()       // Option<&Vec<String>>
            .map(|params| &params.required )
            .into_iter()    // Iterator<&Vec<String>>
            .flatten()      // flattens Iterator<&Vec<String> into Iterator<&String>
            .map(|prop| {
                function.arguments
                    .get(prop)                  // gets the argument value
                    .map(|arg| arg.to_string())
                    .unwrap_or_default()
            })
            .collect();                         // collects the results into a Vec<String>
        (tool.function)(appState, &argv).await
    } else {
        Ok(None)
    }
}

That concludes our toolbox definition. Save and exit the file.

handlers

Edit src/handlers.rs:

server$ vi src/handlers.rs

imports

First modify the use imports at the top of the file:

struct

Next update the following structs:

For the /weather testing API, add also the following struct:

#[derive(Deserialize)]
pub struct Location {
    lat: String,
    lon: String,
}

weather

Let’s implement the handler for the /weather API that we can use to test our getWeather() function later:

pub async fn weather(
    State(appState): State<AppState>,
    ConnectInfo(clientIP): ConnectInfo<&SocketAddr>,
    Json(loc): Json<Location>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let weather = getWeather(
        &appState,
        &vec![loc.lat.to_string(), loc.lon.to_string()],
    )
    .await;
    match weather {
        Err(err) => Err(logServerErr(clientIP, err)),
        Ok(temperature) => {
            logOk(clientIP);
            Ok(Json(json!(temperature)))
        }
    }
}

llmtools

The underlying request/response handling of llmtools() is basically that of llmchat(), plus the mods needed to support tool calling. We will name variables according to this scheme:

Make a copy of your llmchat() function and rename it llmtools(). In your newly renamed llmtools() function, after obtaining a handle to the pgpool, serialize any tools present in the OllamaRequest so that we can add them to the PostgreSQL database:

    // convert tools from client as JSON string (client_tools) and save to db
    let mut client_tools = String::new();
    if let Some(clientTools) = ollamaRequest.tools {
        client_tools = to_string(&clientTools).unwrap_or_default();
    }

Next, when inserting each message into the database, store the client’s tools also, but if there are more than one messages in the messages array, store the tools only once, with the first message. Replace chatterDB.execute("INSERT...) in the for msg in ollamaRequest.messages block with the following:

        chatterDB.execute(
                "INSERT INTO chatts (name, message, id, appID, toolschemas) \
                VALUES ($1, $2, gen_random_uuid(), $3, $4)",
                &[&msg.role, &msg.content, &ollamaRequest.appID, &client_tools], //
            )
            .await
            .map_err(|err| logClientErr(clientIP, StatusCode::NOT_ACCEPTABLE, err.to_string()))?;
        
        // store client_tools only once
        // reset it to empty after first message.
        client_tools = String::new();

The llmchat() code next reconstructs ollamaRequest to be sent to Ollama by retrieving from the PostgreSQL database all prior exchanges between the client and Ollama using the client’s appID. In llmtools(), we first populate ollamRequest.tools with tools resident on the chatterd back end before reconstructing the ollamaRequest. Add before the // reconstruct ollamaRequest to be sent to Ollama: comment:

    // append all chatterd's resident tools to ollamaRequest.tools;
    // front-end tools will be added back later, as part of reconstructing
    // the appID's context from the db (see OllamaMessage.fromRow())
    ollamaRequest.tools = Some(
        TOOLBOX
            .values()
            .map(|tool| tool.schema.clone())
            .collect::<Vec<OllamaToolSchema>>(),
    );

As the comments above indicated, we will need a fromRow() method added to your OllamaMessage struct. Put the fromRow static method for OllamaMessage outside your llmtools() function, for example right under, and also outside, the definition of pub struct OllamaMessage{}, at the top of the file:

impl OllamaMessage {
    fn fromRow(row: &Row, reqTools: &mut Option<Vec<OllamaToolSchema>>) -> OllamaMessage {
        if let Some(toolschemas) = row.get::<usize, Option<String>>(3)
            && let Some(clientTools) = // must deserialize to type to append toolcalls
                from_slice::<Option<Vec<OllamaToolSchema>>>(toolschemas.as_ref()).unwrap_or(None)
        {
            *reqTools = Some(match reqTools.clone() {
                Some(mut tools) => {
                    // append tools to ollama_request.tools
                    tools.extend(clientTools);
                    tools
                }
                None => clientTools,
            })
        }

        OllamaMessage {
            role: row.get(0),
            content: row.get(1),
            toolCalls: if let Some(toolcalls) = row.get::<usize, Option<String>>(2) {
                // has front-end device tools
                // must deserialize to type to append device tools to ollamaRequest.tools         
                from_slice::<std::option::Option<Vec<OllamaToolCall>>>(toolcalls.as_ref())
                    .unwrap_or(None)
            } else {
                None
            },
        }
    }
}

The method fromRow() appends any tools the front-end provided to the ollamaRequest.tools array. This array has previously been populated with available resident back-end tools. Then it creates and returns an OllamaMessage to store a previous exchange between the client and Ollama stored in the given row, including any tool calls Ollama has made.

Back in llmtools(), replace the following lines:

    ollamaRequest.messages = chatterDB
        .query(
            "SELECT name, message FROM chatts WHERE appID = $1 ORDER BY time ASC",
            &[&ollamaRequest.appID],
        )
        .await
        .map_err(|err| logServerErr(&clientIP, err.to_string()))?
        .into_iter()
        .map(|row| OllamaMessage {
            role: row.get(0),
            content: row.get(1),
            toolCalls: None
        })
        .collect();

with:

    ollamaRequest.messages = chatterDB
        .query(
            "SELECT name, message, toolcalls, toolschemas FROM chatts WHERE appID = $1 ORDER BY time ASC",
            &[&ollamaRequest.appID],
        )
        .await
        .map_err(|err| logServerErr(&clientIP, err.to_string()))?
        .into_iter()
        .map(|row| OllamaMessage::fromRow(&row, &mut ollamaRequest.tools))
        .collect();

ndjson_yield_sse

To accommodate resident-tool call, we use a flag, sendNewPrompt, to indicate to our stream generator whether we have any prompt to send to Ollama. Initially, sendNewPrompt is set to true to always send the prompt from the front end. Subsequently, if Ollama makes a call for a tool resident on the back end, we will send the result of the tool call as a new prompt to Ollama. At the top of your let ndjson_yield_sse = stream! { definition, add:

        let mut sendNewPrompt = true;
        let mut tool_result: String;

After we get a handle on pgpool, put the code in the Ok(chatterDB) => {} block inside this while loop, so the new while loop itself is inside the Ok(chatterDB) => {} block:

                while sendNewPrompt {
                    sendNewPrompt = false;   // assume no resident tool call
                    
                    // leave existing code inside `Ok(chatterDB) => {}` block here
                    
                } // while sendNewPrompt                    

Once we get a response from Ollama, add:

                            completion.clear();

Whereas previously in llmchat() we simply yielded each data line after appending it to the completion string, we now must check whether there’s a tool call and yield the data line only if there were no tool call. Replace the following lines:

                            // send NDJSON line as SSE data line
                            yield Ok(Event::default().data(&line));

with:

                                // is there a tool call?
                                if let Some(toolCalls) = &ollamaResponse.message.toolCalls {
                                    // handle tool calls

                                } else {
                                    // no tool call, send NDJSON line as SSE data line
                                    yield Ok(Event::default().data(&line));
                                }

In handling tool calls, we first serialize the tool call back into a JSON string to be saved into the database. Replace the comment // handle tool calls with:

                                    // convert toolCalls to JSON string (tool_calls) to be saved to db
                                    let mut tool_calls = to_string(&toolCalls).unwrap_or_default();

                                    for toolCall in toolCalls {
                                        // but assuming one tool call per response
                                        if toolCall.function.name.is_empty() {
                                            continue // LLM miscalled
                                        }

                                        // save full response, including tool call(s), to db,
                                        // to form part of session's prompt history
                                        if let Err(err) = chatterDB.execute(
                                            "INSERT INTO chatts (name, message, id, appid, toolcalls) \
                                                VALUES ('assistant', $1, gen_random_uuid(), $2, $3)",
                                            &[&completion, &ollamaRequest.appID, &tool_calls],
                                        ) .await {
                                            yield Ok(Event::default().event("error")
                                                .data(json!({ "error": err.to_string() }).to_string()))
                                        };
                                        
                                        // clear completion and tool_calls, we already stored them
                                        completion.clear();
                                        tool_calls.clear();
                                        
                                        // make the tool call

We call toolInvoke() with the tool’s signature and process the result. There are three possible outcomes from the call to toolInvoke():

  1. the tool is resident but the call was unsuccesfull and returned an error,
  2. the tool is resident and the call was successful, or
  3. the tool is non-resident.

If the tool call resulted in an error, we store the error as the tool result. We add the tool call and its result to the OllamaRequest message and set the flag (sendNewPrompt) to send the OllamaRequest back to Ollama. We also store both the tool call and its result to the database, to form part of this appID’s context. If the tool call resulted in neither an error nor any returned result, we interpret that as the tool being non-resident on the back end and forward the tool call to the front end as an SSE tool_calls event. Replace the comment # make the tool call with:

                                        let toolResult = toolInvoke(&appState, toolCall.clone().function).await;
                                        match toolResult {
                                            Err(err) => {
                                                // outcome 1: tool resident but had error
                                                // send error back to LLM, don't report to frontend
                                                tool_result = err;
                                            },
                                            Ok(result) => {
                                                // outcome 2: tool call is resident and no error
                                                tool_result = result.unwrap_or_default();
                                            }
                                        }

                                        if tool_result.is_empty() {
                                            // outcome 3: tool non resident, forward to
                                            // front end as 'tool_calls' SSE event
                                            yield Ok(Event::default().event("tool_calls")
                                                .data(&line))
                                        } else {
                                            // reuse OllamaMessage to carry tool result
                                            // to be sent back to Ollama:
                                            // first append the tool call itself
                                            ollamaRequest.messages.push(ollamaResponse.message.to_owned());
                                            // then append the result
                                            ollamaRequest.messages.push(OllamaMessage {
                                                role: "tool".to_string(),
                                                content: tool_result.to_owned(),
                                                toolCalls: None,
                                            });

                                            // don't send tools multiple times
                                            ollamaRequest.tools = None;
                                            // loop to send tool result back to Ollama
                                            sendNewPrompt = true;

                                            // save resident tool call result or error message
                                            if let Err(err) = chatterDB.execute(
                                                "INSERT INTO chatts (name, message, id, appid) \
                                                VALUES ('tool', $1, gen_random_uuid(), $2)",
                                                &[&wsRegex.replace_all(&to_string(&*tool_result)
                                                .unwrap_or_default(), " ").to_string(), &ollamaRequest.appID],
                                            ) .await {
                                                yield Ok(Event::default().event("error")
                                                    .data(json!({ "error": err.to_string() }).to_string()))
                                            };
                                        }
                                    } // for toolCall
                                    

We keep the rest of the llmchat() code without further changes and we’re done with handlers.rs! Save and exit the file.

main.rs

Edit src/main.rs:

server$ vi src/main.rs

Find mod handlers; and add below it:

mod toolbox;

Find the router = Router::new() instantiation statement and add these routes right after the route for /llmchat:

        .route("/llmtools", post(handlers::llmtools))
        .route("/weather", get(handlers::weather))

We’re done with main.rs. Save and exit the file.

Build and test run

To build your server:

server$ cargo build --release

As before, it will take some time to download and build all the 3rd-party crates. Be patient.

Linking error with cargo build?

When running cargo build --release, if you see:

  error: linking with cc failed: exit status: 1
  note: collect2: fatal error: ld terminated with signal 9 [Killed]

below a long list of object files, try running cargo build --release again. It usually works the second time around, when it will have less remaining linking to do. If the error persisted, please talk to the teaching staff.


:point_right:Rust is a compiled language, like C/C++ and unlike Python, which is an interpreted language. This means you must run cargo build each and every time you made changes to your code, for the changes to show up in your executable.

To run your server:

server$ sudo ./chatterd
# Hit ^C to end the test

You can test your implementation following the instructions in the Testing llmTools APIs section.


Prepared by Xin Jie ‘Joyce’ Liu, Chenglin Li, and Sugih Jamin Last updated March 9th, 2026