Originally published on tamiz.pro.
The prevailing narrative in artificial intelligence has been dominated by cloud-based, API-driven models. While this approach offers scalability, it introduces critical latency, dependency on external services, and significant privacy concerns regarding data exfiltration. For mission-critical applications, financial analysis, or healthcare systems, the inability to guarantee data residency and offline operation is a non-starter.
The solution lies in a "Local-First" architecture, where the AI assistant runs entirely on-premise or on-device. However, building such systems requires more than just down an LLM weights file; it demands a robust infrastructure layer capable of managing state, memory safety, and real-time concurrency. This article explores how to construct this infrastructure using two powerhouse languages: Go for its superior concurrency primitives and developer velocity in orchestration, and Rust for its memory safety, zero-cost abstractions, and performance-critical inference execution.
We will dissect the architecture of a secure, local-first AI agent, moving from the conceptual model to the implementation details, focusing on the boundary between the orchestration layer (Go) and the execution layer (Rust).
Building a local-first AI assistant is not merely a software engineering challenge; it is a systems architecture problem. The core tension lies between flexibility (the ability to swap models, adjust prompts, and handle complex workflows) and performance/security (minimizing latency and preventing memory corruption or data leaks).
To resolve this, we adopt a micro-kernel architecture:
This separation allows teams to iterate rapidly on the agent's logic in Go (leveraging its vast ecosystem for HTTP servers, database drivers, and UI frameworks) while maintaining a hardened, performant core in Rust.
Local-first AI is computationally intensive. Unlike cloud inference, where you can scale horizontally indefinitely, local inference is bound by the hardware constraints of the host machine (CPU/RAM/GPU). Rust is chosen for the engine layer for three primary reasons:
AI models often process unstructured text, which can be adversarial. Malformed inputs can lead to buffer overflows in C/C++ libraries. Rust’s ownership model guarantees memory safety at compile time. For a local-first agent, this is a security feature, not just a performance one. If the inference engine crashes due to a memory error, it takes down the entire local service. Rust prevents this.
WebAssembly (Wasm) and real-time systems require deterministic behavior. Rust’s lack of garbage collection s ensures that token generation latency remains consistent, which is crucial for chat interfaces that expect real-time streaming responses.
The modern Rust ML ecosystem, including Burn, Candle (by Hugging Face), and tch-rs (PyTorch bindings), provides access to state-of-the-art models. These libraries are optimized for SIMD instructions and multi-threaded matrix multiplication, squeezing every bit of performance out of consumer hardware.
While Rust is superior for raw computation, Go excels at concurrency patterns and ecosystem integration. An AI agent is rarely just a model; it is a system that:
Go’s goroutine
model is perfectly suited for this. Each user session can be assigned a dedicated goroutine, allowing the system to handle thousands of concurrent users with a small memory footprint. Furthermore, Go’s standard library provides excellent tools for building HTTP/2 servers, handling WebSocket streams for real-time token delivery, and interacting with SQL/NoSQL databases.
Let’s visualize the data flow in a typical request:
The cost of inter-process communication (IPC) or Foreign Function Interface (FFI) calls can be significant if data is copied excessively. To mitigate this, we use zero-copy strategies where possible.
For gRPC, we define a schema that minimizes payload size. For FFI (Go calling Rust), we use unsafe
blocks carefully to pass pointers to pre-allocated buffers, avoiding the serialization/deserialization overhead of JSON.
We will use Candle (by Hugging Face) for its simplicity and performance in local inference. Candle is designed to be embeddable and works well on CPU and GPU.
rust-engine/
├── Cargo.toml
├── src/
│ ├── main.rs # Entry point, gRPC server
│ ├── inference.rs # Model and execution logic
│ ├── tokenizer.rs # Tokenization utilities
│ └── lib.rs # Exports for FFI/gRPC
└── proto/
└── agent.proto # gRPC definitions
First, we define the contract between Go and Rust. This ensures type safety and clear expectations.
// proto/agent.proto
syntax = "proto3";
package agent;
service InferenceService {
// Stream tokens to the orchestrator
rpc GenerateStream (Request) returns (stream Response);
// Health check and model status
rpc GetModelStatus (Empty) returns (Status);
}
message Request {
string prompt = 1;
int32 max_tokens = 2;
float temperature = 3;
repeated string history = 4; // Previous conversation turns
}
message Response {
string token = 1;
bool is_end = 2;
}
message Status {
string model_name = 1;
bool is_ready = 2;
}
message Empty {}
We use tonic
for gRPC and candle
for inference. This setup allows the Rust binary to act as a standalone service that Go clients connect to.
// src/inference.rs
use candle::{Device, Tensor, DType};
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::llama as llama_model;
use std::sync::Arc;
pub struct LlamaEngine {
model: Arc<llama_model::Llama>,
device: Device,
tokenizer: tokenizers::Tokenizer,
}
impl LlamaEngine {
pub fn new(model_path: &str, tokenizer_path: &str, device: Device) -> Result<Self, Box<dyn std::error::Error>> {
// Load tokenizer
let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)?;
// Load model weights
let model = llama_model::Llama::from_preferred(model_path, &device)?;
Ok(LlamaEngine {
model: Arc::new(model),
device,
tokenizer,
})
}
pub fn generate(&self, prompt: &str, max_tokens: usize, temperature: f64) -> Result<String, Box<dyn std::error::Error>> {
// Tokenize input
let tokens = self.tokenizer.encode(prompt, true)?;
let tokens = tokens.get_ids();
// Prepare input tensor
let input = Tensor::new(tokens, &self.device)?;
// Initialize logits processor
let mut logits_processor = LogitsProcessor::new(temperature, Some(42), None);
let mut output_tokens = Vec::new();
let mut current_tokens = input.clone();
// Inference loop
for _ in 0..max_tokens {
let logits = self.model.forward(¤t_tokens, 0)?;
let logits = logits.squeeze(0)?;
let next_token = logits_processor.sample(&logits)?;
output_tokens.push(next_token);
// Break if end of sequence
if next_token == self.tokenizer.tokenizer.eos_token_id() {
break;
}
// Prepare next input
current_tokens = Tensor::new(&[next_token], &self.device)?;
}
// Decode output tokens
let decoded = self.tokenizer.decode(&output_tokens, true)?;
Ok(decoded)
}
}
// src/main.rs
use tonic::{transport::Server, Request, Response, Status};
use tokio_stream::{Stream, StreamExt};
mod inference;
mod agent {
include!(concat!(env!("OUT_DIR"), "/agent.rs"));
}
struct InferenceServiceImpl {
engine: inference::LlamaEngine,
}
#[tonic::async_trait]
impl agent::inference_service_server::InferenceService for InferenceServiceImpl {
async fn generate_stream(
&self,
request: Request<agent::Request>,
) -> Result<Response<Self::GenerateStreamStream>, Status> {
let req = request.into_inner();
let prompt = req.prompt;
let max_tokens = req.max_tokens as usize;
let temperature = req.temperature;
// Generate tokens synchronously for simplicity,
// but in production, you'd stream intermediate tokens
let result = match self.engine.generate(&prompt, max_tokens, temperature) {
Ok(text) => text,
Err(e) => return Err(Status::internal(e.to_string())),
};
// For streaming, we'd split the result and yield chunks
// Here we return a single response for demonstration
Ok(Response::new(agent::Response {
token: result,
is_end: true,
}))
}
async fn get_model_status(
&self,
_request: Request<agent::Empty>,
) -> Result<Response<agent::Status>, Status> {
Ok(Response::new(agent::Status {
model_name: "Llama-2-7b".to_string(),
is_ready: true,
}))
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let engine = inference::LlamaEngine::new(
"./models/llama-2-7b",
"./models/tokenizer.json",
Device::Cpu,
)?;
let svc = agent::inference_service_server::InferenceServiceServer::new(
InferenceServiceImpl { engine }
);
Server::builder()
.add_service(svc)
.serve("127.0.0.1:50051".parse()?)
.await?;
Ok(())
}
The Go application acts as the brain, managing the state of the conversation and coordinating with the Rust engine.
We need a way to manage multiple users. A simple in-memory map is sufficient for demonstration, but in production, you’d use Redis or PostgreSQL.
// go-orchestrator/internal/session/store.go
package session
import (
"sync"
)
type Session struct {
ID string
History []string
}
type Store struct {
mu sync.RWMutex
sessions map[string]*Session
}
func NewStore() *Store {
return &Store{
sessions: make(map[string]*Session),
}
}
func (s *Store) GetOrCreate(id string) *Session {
s.mu.RLock()
sess, ok := s.sessions[id]
s.mu.RUnlock()
if !ok {
s.mu.Lock()
// Double-check locking
if sess, ok = s.sessions[id]; !ok {
sess = &Session{ID: id, History: make([]string, 0)}
s.sessions[id] = sess
}
s.mu.Unlock()
}
return sess
}
The Go app needs to connect to the Rust gRPC server. We use the generated protobuf client.
// go-orchestrator/internal/ai/engine.go
package ai
import (
"context"
"log"
"github.com/yourname/agent/proto/agent"
"google.golang.org/grpc"
)
type RustEngineClient struct {
client agent.InferenceServiceClient
}
func NewRustEngineClient(addr string) (*RustEngineClient, error) {
conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
return nil, err
}
return &RustEngineClient{
client: agent.NewInferenceServiceClient(conn),
}, nil
}
func (c *RustEngineClient) Generate(ctx context.Context, prompt string, history []string, maxTokens int, temperature float32) (string, error) {
resp, err := c.client.GenerateStream(ctx, &agent.Request{
Prompt: prompt,
MaxTokens: int32(maxTokens),
Temperature: temperature,
History: history,
})
if err != nil {
return "", err
}
return resp.Token, nil
}
Go’s strength is handling many concurrent connections. We use a WebSocket server to stream responses back to the client.
// go-orchestrator/internal/handler/ai_handler.go
package handler
import (
"context"
"log"
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/yourname/agent/go-orchestrator/internal/ai"
"github.com/yourname/agent/go-orchestrator/internal/session"
)
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
func AIHandler(engine *ai.RustEngineClient, store *session.Store) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println("Upgrade error:", err)
return
}
defer conn.Close()
// Generate a unique session ID for this connection
sessionID := "user-" + fmt.Sprintf("%d", time.Now().UnixNano())
sess := store.GetOrCreate(sessionID)
ctx := context.Background()
for {
_, msg, err := conn.ReadMessage()
if err != nil {
break
}
// Process message
response, err := engine.Generate(ctx, string(msg), sess.History, 500, 0.7)
if err != nil {
conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
continue
}
// Send response
conn.WriteMessage(websocket.TextMessage, []byte(response))
// Update history
sess.History = append(sess.History, string(msg), response)
// Limit history size to manage memory
if len(sess.History) > 20 {
sess.History = sess.History[len(sess.History)-20:]
}
}
}
}
Building a local-first AI assistant introduces unique security challenges. Unlike cloud APIs, where you can implement rate limiting and WAFs, your local application is the first line of defense.
Even locally, users can attempt prompt injection. The Go orchestrator should sanitize inputs before sending them to the Rust engine. Implement a pre-processing layer that filters out malicious patterns or restricts the context window to prevent buffer overflow attempts in the Rust layer.
If the Rust engine is compiled as a shared library (.so
or .dylib
) and linked into the Go binary, a vulnerability in Rust could compromise the entire process. To mitigate this, consider running the Rust engine as a separate process (as shown in the gRPC example) and communicating via IPC. This way, a crash in the Rust engine does not take down the Go orchestrator, and you can implement stricter OS-level sandboxing for the Rust process.
Local-first implies data stays on the device. Ensure that session histories and any sensitive data processed by the AI are encrypted at rest. Use Go’s crypto
package to encrypt the session store before writing to disk.
Local inference is constrained by hardware. To make the assistant responsive, you must optimize the model.
Large Language Models are often stored in FP16 or BF16. Quantizing to INT8 or even INT4 can reduce memory usage by 2-4x with minimal loss in quality. The Candle library supports quantized models out of the box. Load quantized models to fit larger context windows on consumer hardware.
If you have multiple users, the Rust engine can batch requests. However, for a local-first single-user assistant, batching is less relevant. For multi-user scenarios, the Rust engine can maintain a batch queue, processing multiple prompts in parallel using multi-threaded matrix multiplication.
Building a secure, local-first AI assistant is a complex but rewarding engineering challenge. By leveraging Go for its concurrency and ecosystem, and Rust for its performance and memory safety, you can create a system that is both flexible and robust.
This architecture provides:
As AI moves from a cloud-centric model to a distributed, edge-centric one, this hybrid approach will become the standard for high-assurance applications. Start by prototyping the Rust inference engine, then build the Go orchestrator around it, and gradually add security and optimization layers.
Q: Can I use Python for the orchestration layer?
A: Yes, Python is often used for AI development. However, for a high-concurrency local service, Go or Rust is superior. Python’s GIL limits true parallelism, making it less ideal for handling many simultaneous WebSocket connections compared to Go’s goroutines.
Q: How do I handle model updates in a local-first app?
A: Implement a versioning system in your gRPC protocol. When a new model is released, the Go orchestrator can trigger a download of the new weights (quantized) and reload the Rust engine process without restarting the entire application.
Q: Is this approach viable for mobile devices?
A: Yes. The same architecture can be adapted for iOS/Android using Go (via gomobile) or Rust (via Swift/Kotlin bindings). The key is to keep the inference engine (Rust) native and the logic (Go) as a thin layer or also native.