smista.ai
smista.ai · blog
How we built a provider-agnostic LLM layer in Rust with Rig
A practical look at using Rig to build a Rust abstraction over multiple LLM providers, tools, and memory.
Every provider has a different client
One of our earliest goals for smista.ai was to support models from multiple providers.
In most projects that interface with LLMs, there is often a need to communicate with multiple providers, such as Anthropic and OpenAI, to meet the user's needs.
Once you support more than one LLM provider, the problem stops being "how do I call an API?" and becomes "how do I avoid leaking provider-specific details everywhere?"
At smista.ai, we wanted an integration that supports at least five providers: Anthropic, OpenAI, Gemini, Ollama, and any OpenAI-compatible endpoint. We wanted this integration to be exposed through a common trait, so the rest of the router would not have to care which provider was behind a request.
This post reflects smista-providers as it stands on rig 0.42 (rig-core
and rig-agent).
Wrapping Rig's CompletionClient
rig gives every provider a CompletionClient that builds a
provider-specific CompletionModel. Our own Agent<C> wraps one of those:
pub struct Agent<C>
where
C: CompletionClient,
{
model: C::CompletionModel,
preamble: String,
internal_tool: Arc<dyn InternalTool>,
internal_tool_definition: RigToolDefinition,
descriptor: ModelDescriptor,
}We deliberately do not use rig's own high-level Agent builder and its built-in
tool-calling loop: we need to tell apart tool calls our own memory tool can
execute on the spot from tool calls the router has to mediate (the model's own
tools), and a loop we don't control can't make that distinction. So Agent only
stores the bare CompletionModel and its own preamble, and builds a fresh
CompletionRequestBuilder for every turn:
fn request_builder(
&self,
prompt: RigMessage,
history: &[RigMessage],
parameters: &ModelParameters,
tools: Vec<ToolDefinition>,
tool_choice: ToolChoice,
) -> CompletionRequestBuilder<C::CompletionModel> {
let mut all_tools: Vec<RigToolDefinition> = tools.into_iter().map(into_rig_tool).collect();
all_tools.push(self.internal_tool_definition.clone());
let builder = self
.model
.completion_request(prompt)
.preamble(self.preamble.clone())
.messages(history.iter().cloned())
.tools(all_tools)
.temperature_opt(parameters.temperature.map(f64::from))
.max_tokens_opt(parameters.max_tokens.map(u64::from))
.additional_params_opt(additional_params(parameters));
match into_rig_tool_choice(tool_choice) {
Some(choice) => builder.tool_choice(choice),
None => builder,
}
}Every request carries the agent's own preamble and its internal tool definition (the memory tool) alongside whatever tools the caller passed in, so the model always sees the same system context regardless of which turn it is on.
From clients to abstract Models
But how do we get from a connection to Gemini or Anthropic to our Agent
struct, and from there to something the router can call?
Rig already provides provider-specific clients that can be configured with the
required settings, such as API keys, base URLs, and other provider-specific
options. That let us keep our own abstraction small: each provider gets a thin
*Model type that owns an Agent<C> and implements our own Model trait.
For example, our Gemini model works like this:
pub struct GeminiModel {
agent: Agent<GeminiClient>,
descriptor: ModelDescriptor,
reference: ModelReference,
}
impl GeminiModel {
pub async fn new<S>(
GeminiModelArgs { preamble, storage }: GeminiModelArgs<S>,
authentication: &Authentication,
descriptor: ModelDescriptor,
scope: MemoryScope,
preamble_segments: &[String],
) -> Result<Self, ProviderError>
where
S: MemoryStorage + 'static,
{
let reference = descriptor.reference();
let api_key = authentication.require_api_key(Provider::Gemini, &reference.model)?;
let client = GeminiClient::new(api_key.expose_secret())?;
let agent = Agent::new(AgentArgs {
completion_model: client,
descriptor: descriptor.clone(),
preamble,
preamble_segments: preamble_segments.to_vec(),
storage,
scope,
})
.await?;
Ok(Self { agent, descriptor, reference })
}
}While the Anthropic and Ollama clients are similar, they each use their own rig
client instead. But all of them implement CompletionClient, which is all
Agent<C> needs to build itself. scope tells the memory backend which user
and session the agent's memory reads and writes belong to, and
preamble_segments lets a caller append extra instructions (skills, per-turn
context) to the preamble without the model layer knowing what they contain.
Sending a request from an Agent to the LLM, one turn at a time
Agent::complete is where the internal/external tool-call split actually
happens. It sends a turn, and if every tool call the model asked for is one our
memory tool can execute, it runs them and immediately sends a follow-up turn
instead of handing the tool calls back to the router:
pub async fn complete(&self, request: CompletionRequest) -> ProviderResult<CompletionResponse> {
let CompletionRequest { messages, parameters, tools, tool_choice } = request;
// ...build `history` and pop the last message as `prompt`...
let mut usage = RigUsage::new();
for _ in 0..MAX_INTERNAL_TOOL_TURNS {
let response = self
.request_builder(prompt.clone(), &history, ¶meters, tools.clone(), tool_choice)
.send()
.await
.map_err(|error| self.error(category_from_completion(&error), error.to_string()))?;
usage += response.usage;
// ...collect text and tool calls from `response.choice`...
let all_internal = !tool_calls.is_empty()
&& tool_calls.iter().all(|call| call.function.name == self.internal_tool_definition.name);
if !all_internal {
// at least one tool call belongs to the router: return as-is
return Ok(CompletionResponse { /* ... */ });
}
// every call this turn is agent-internal: execute them and loop
let mut results = Vec::with_capacity(tool_calls.len());
for call in &tool_calls {
results.push(self.execute_internal_tool(call).await?);
}
history.push(prompt);
history.push(RigMessage::Assistant { id: response.message_id.clone(), content: response.choice.clone() });
prompt = RigMessage::User { content: results };
}
Err(self.error(ProviderErrorCategory::Unknown, "model exceeded the internal tool turn limit"))
}MAX_INTERNAL_TOOL_TURNS exists only to stop a model that keeps calling the
memory tool from looping forever; a normal memory operation resolves in one or
two turns.
stream() is the streaming counterpart: it calls .stream() instead of
.send() and maps rig's StreamedAssistantContent items onto our own
StreamEvent vocabulary (text deltas, tool-call start/requested, usage, done).
Unlike complete(), streaming never executes a tool itself — every call,
internal or not, comes out as a StreamEvent::ToolCallRequested and the caller
decides what to do with it. Models without the streaming capability fall back
to complete() and have its single response replayed as a short stream.
This is, of course, a trimmed-down version of the real thing. See agent.rs on smista.ai repository for the full implementation, including how text-only tool calls from providers that don't emit structured ones get recovered, and how usage and pricing are computed.
Adding user and session memory through a Rig PortableTool
One of the core functionalities of LLM agents is the possibility to use session and user memories. This means we must be able to save meaningful information for reuse in the current session and for the long term, including user preferences. Later, our agent must be capable of recalling memories for use.
For this purpose, we have implemented a MemoryStorage trait to provide this
functionality and a rig PortableTool.
To store session and user data, we have defined the MemoryStorage trait,
which, on smista.ai, is later implemented as a client for SurrealDB. The backend
is a single long-lived, shared handle rather than something constructed already
scoped to one user: every call carries a MemoryScope (a user_id and a
session_id) telling the backend which owner it is acting for.
First, we have defined a MemoryRecord as an entity to keep track of memories
with a key to uniquely identify each piece of information:
pub struct MemoryRecord {
/// Opaque, backend-defined handle identifying this entry. Pass it back to
/// `forget_*` to remove exactly this entry.
pub handle: String,
/// Optional topic. A keyed entry upserts on its key, so a later record with
/// the same key replaces this one; keyless entries accumulate freely.
pub key: Option<String>,
/// The remembered fact.
pub content: String,
}Then we have used this data as the core type for the MemoryStorage:
pub trait MemoryStorage: Send + Sync {
type Error: std::error::Error + Send + Sync + 'static;
fn put_user_memory(
&self,
scope: MemoryScope,
key: Option<String>,
content: String,
) -> impl Future<Output = Result<MemoryRecord, Self::Error>> + Send;
fn forget_user_memory(
&self,
scope: MemoryScope,
handle: String,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn get_user_memories(
&self,
scope: MemoryScope,
limit: Option<usize>,
) -> impl Future<Output = Result<Vec<MemoryRecord>, Self::Error>> + Send;
fn get_user_memory_by_key(
&self,
scope: MemoryScope,
key: String,
) -> impl Future<Output = Result<Option<MemoryRecord>, Self::Error>> + Send;
}The real trait then exposes the same methods for session records as well.
To collect operations on user memories, we need a tool. A tool is a typed function exposed to the model. The model does not execute it directly; it emits a structured tool call, and the host application decides whether and how to run it.
rig exposes this as the PortableTool trait: description() and
parameters() are plain synchronous methods, and the tool's definition is
derived from them rather than built by hand:
pub struct MemoryTool<S>
where
S: MemoryStorage,
{
/// The shared backend this tool records into and forgets from.
storage: Arc<S>,
/// The user and session this tool's operations are scoped to.
scope: MemoryScope,
}
impl<S> PortableTool for MemoryTool<S>
where
S: MemoryStorage + 'static,
{
const NAME: &'static str = "memory";
type Args = MemoryArgs;
type Output = String;
type Error = MemoryToolError<S::Error>;
fn description(&self) -> String {
concat!(
"Record or forget a memory addressed by `key`.\n\n",
"Use scope `user` for durable facts about the user that should ",
"persist across sessions (preferences, identity, long-lived ",
"context). Use scope `session` for working memory tied to the ",
"current session only.\n\n",
"Operations:\n",
"- `record`: store `value` under `key`, replacing any existing ",
"fact filed under the same key.\n",
"- `forget`: remove the fact filed under `key`.\n\n",
"You do not need to recall memories: everything recorded is ",
"already provided to you as context at the start of the turn."
)
.to_string()
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"op": {
"type": "string",
"enum": ["record", "forget"],
"description": "The operation to perform."
},
"scope": {
"type": "string",
"enum": ["user", "session"],
"description": "Which store to target: durable user memory or session-only working memory."
},
"key": {
"type": "string",
"description": "Topic the fact is filed under; reuse the same key to replace or forget it."
},
"value": {
"type": "string",
"description": "The fact to record. Required for `record`, ignored for `forget`."
}
},
"required": ["op", "scope", "key"]
})
}
async fn call(&self, args: MemoryArgs) -> Result<String, Self::Error> {
let MemoryArgs { op, scope: store, key, value } = args;
match op {
MemoryOp::Record => {
let value = value.ok_or(MemoryToolError::MissingValue)?;
match store {
MemoryStore::User => {
self.storage.put_user_memory(self.scope, Some(key.clone()), value).await?;
}
MemoryStore::Session => {
self.storage.put_session_memory(self.scope, Some(key.clone()), value).await?;
}
}
Ok(format!("Recorded {} memory \"{key}\".", store.label()))
}
MemoryOp::Forget => match self.handle_for(store, key.clone()).await? {
Some(handle) => {
match store {
MemoryStore::User => self.storage.forget_user_memory(self.scope, handle).await?,
MemoryStore::Session => self.storage.forget_session_memory(self.scope, handle).await?,
}
Ok(format!("Forgot {} memory \"{key}\".", store.label()))
}
None => Ok(format!("No {} memory found for \"{key}\".", store.label())),
},
}
}
}Note the two distinct notions of "scope" here. self.scope: MemoryScope is
who the tool is acting for (the user and session it was built for), fixed at
construction. MemoryStore (the op's scope argument, user or session)
is which store the model wants to touch, chosen per call. Keeping them as two
separate types stops the two ideas from getting conflated.
description and parameters tell the LLM when the tool should be called, what
it does, and how to pass data to it; call is the function executed whenever
the model triggers it.
Notice that the tool does not expose a recall operation. This is intentional: memory retrieval is handled by the host before the turn starts, and the resulting memories are appended to the agent preamble. The model can record or forget memories, but it does not decide how memory retrieval works.
Linking memory into the Agent
At this point, we can pass our MemoryStorage implementation to Agent::new
and link both the tool and the user's records to it. We keep the memory tool
ourselves — internal_tool on Agent<C> — rather than handing it off to rig,
so complete() can decide whether to run it or forward it to the router:
pub async fn new<S>(
AgentArgs { completion_model, descriptor, preamble, preamble_segments, storage, scope }: AgentArgs<C, S>,
) -> ProviderResult<Self>
where
S: MemoryStorage + 'static,
{
let model_name = descriptor.model.clone();
let provider = descriptor.provider.clone();
// load preamble from memory storage, scoped to this user and session
let memory_preamble =
load_memories_preamble(storage.as_ref(), scope, provider.clone(), &model_name).await?;
// load memory tool, scoped the same way
let memory_tool = MemoryTool::new(storage.clone(), scope);
// `preamble` replaces the system prompt, so the memory preamble and any
// extra segments must be appended rather than set, or they would wipe
// the base preamble. Final order: base, then memory, then segments.
let mut full_preamble = preamble;
if let Some(memories) = memory_preamble {
full_preamble.push('\n');
full_preamble.push_str(&memories);
}
for segment in &preamble_segments {
full_preamble.push('\n');
full_preamble.push_str(segment);
}
let internal_tool_definition = InternalTool::definition(&memory_tool);
let model = completion_model.completion_model(model_name.clone());
Ok(Self {
model,
preamble: full_preamble,
internal_tool: Arc::new(memory_tool),
internal_tool_definition,
descriptor,
})
}InternalTool is a small trait we erase every PortableTool we execute
ourselves through, so Agent<C> can stay generic only over the completion
client and not over every tool's own generic parameters — the memory tool is
itself generic over its MemoryStorage backend.
Conclusion
Rig did not remove the need for our own abstraction, but it moved the provider-specific complexity to the edges.
Each provider still has its own client, authentication, model names, and
configuration. But once a client implements CompletionClient, the rest of
smista.ai can build the same Agent, attach the same internal tool, inject the
same memory preamble, and handle completions through the same flow.
Owning the turn loop ourselves, instead of rig's built-in one, is what lets us
treat an internal tool call (memory) and a router-mediated one differently.
PortableTool keeps a tool's definition and its logic in one place without an
async call just to describe itself.
That is exactly the kind of boundary we wanted: provider-specific where necessary, provider-agnostic everywhere else.
You can find the full implementation in the smista-providers crate on the smista.ai repository.
Reference: