refactor RemoteHttpCompute config

seun-ja · seun-ja · commit a2d3bc2da13a · 2025-08-28T22:45:33.000+01:00
CustomLlm is parsed directly, catches unsupported CustomLlm. Also introduces a new trait called LlmWorker which every LLM engine implements

Signed-off-by: Aminu Oluwaseun Joshua &lt;seun.aminujoshua@gmail.com&gt;
diff --git a/crates/factor-llm/src/spin.rs b/crates/factor-llm/src/spin.rs
@@ -2,7 +2,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use spin_factors::runtime_config::toml::GetTomlValue;
-use spin_llm_remote_http::RemoteHttpLlmEngine;
+use spin_llm_remote_http::{CustomLlm, RemoteHttpLlmEngine};
 use spin_world::async_trait;
 use spin_world::v1::llm::{self as v1};
 use spin_world::v2::llm::{self as v2};
@@ -122,7 +122,7 @@ impl LlmCompute {
             LlmCompute::RemoteHttp(config) => Arc::new(Mutex::new(RemoteHttpLlmEngine::new(
                 config.url,
                 config.auth_token,
-                config.custom_llm.and_then(|c| c.as_str().try_into().ok()),
+                config.custom_llm,
             ))),
         };
         Ok(engine)
@@ -133,7 +133,8 @@ impl LlmCompute {
 pub struct RemoteHttpCompute {
     url: Url,
     auth_token: String,
-    custom_llm: Option<String>,
+    #[serde(default)]
+    custom_llm: CustomLlm,
 }
 
 /// A noop engine used when the local engine feature is disabled.
diff --git a/crates/llm-remote-http/src/default.rs b/crates/llm-remote-http/src/default.rs
@@ -4,27 +4,43 @@ use reqwest::{
     Client, Url,
 };
 use serde_json::json;
-use spin_world::v2::llm::{self as wasi_llm};
+use spin_world::{
+    async_trait,
+    v2::llm::{self as wasi_llm},
+};
 
-use crate::{EmbeddingResponseBody, InferRequestBodyParams, InferResponseBody};
+use crate::{EmbeddingResponseBody, InferRequestBodyParams, InferResponseBody, LlmWorker};
 
-pub(crate) struct DefaultAgentEngine;
+pub(crate) struct DefaultAgentEngine {
+    auth_token: String,
+    url: Url,
+    client: Option<Client>,
+}
 
 impl DefaultAgentEngine {
-    pub async fn infer(
-        auth_token: &String,
-        url: &Url,
-        mut client: Option<Client>,
+    pub fn new(auth_token: String, url: Url, client: Option<Client>) -> Self {
+        Self {
+            auth_token,
+            url,
+            client,
+        }
+    }
+}
+
+#[async_trait]
+impl LlmWorker for DefaultAgentEngine {
+    async fn infer(
+        &mut self,
         model: wasi_llm::InferencingModel,
         prompt: String,
         params: wasi_llm::InferencingParams,
     ) -> Result<wasi_llm::InferencingResult, wasi_llm::Error> {
-        let client = client.get_or_insert_with(Default::default);
+        let client = self.client.get_or_insert_with(Default::default);
 
         let mut headers = HeaderMap::new();
         headers.insert(
             "authorization",
-            HeaderValue::from_str(&format!("bearer {}", auth_token)).map_err(|_| {
+            HeaderValue::from_str(&format!("bearer {}", self.auth_token)).map_err(|_| {
                 wasi_llm::Error::RuntimeError("Failed to create authorization header".to_string())
             })?,
         );
@@ -45,7 +61,8 @@ impl DefaultAgentEngine {
         }))
         .map_err(|_| wasi_llm::Error::RuntimeError("Failed to serialize JSON".to_string()))?;
 
-        let infer_url = url
+        let infer_url = self
+            .url
             .join("/infer")
             .map_err(|_| wasi_llm::Error::RuntimeError("Failed to create URL".to_string()))?;
         tracing::info!("Sending remote inference request to {infer_url}");
@@ -68,19 +85,17 @@ impl DefaultAgentEngine {
         }
     }
 
-    pub async fn generate_embeddings(
-        auth_token: &String,
-        url: &Url,
-        mut client: Option<Client>,
+    async fn generate_embeddings(
+        &mut self,
         model: wasi_llm::EmbeddingModel,
         data: Vec<String>,
     ) -> Result<wasi_llm::EmbeddingsResult, wasi_llm::Error> {
-        let client = client.get_or_insert_with(Default::default);
+        let client = self.client.get_or_insert_with(Default::default);
 
         let mut headers = HeaderMap::new();
         headers.insert(
             "authorization",
-            HeaderValue::from_str(&format!("bearer {}", auth_token)).map_err(|_| {
+            HeaderValue::from_str(&format!("bearer {}", self.auth_token)).map_err(|_| {
                 wasi_llm::Error::RuntimeError("Failed to create authorization header".to_string())
             })?,
         );
@@ -95,7 +110,7 @@ impl DefaultAgentEngine {
         let resp = client
             .request(
                 reqwest::Method::POST,
-                url.join("/embed").map_err(|_| {
+                self.url.join("/embed").map_err(|_| {
                     wasi_llm::Error::RuntimeError("Failed to create URL".to_string())
                 })?,
             )
@@ -114,4 +129,8 @@ impl DefaultAgentEngine {
             ))),
         }
     }
+
+    fn url(&self) -> Url {
+        self.url.clone()
+    }
 }
diff --git a/crates/llm-remote-http/src/lib.rs b/crates/llm-remote-http/src/lib.rs
@@ -1,53 +1,47 @@
 use anyhow::Result;
-use reqwest::{Client, Url};
+use reqwest::Url;
 use serde::{Deserialize, Serialize};
-use spin_world::v2::llm::{self as wasi_llm};
-
-use crate::{
-    default::DefaultAgentEngine,
-    open_ai::OpenAIAgentEngine,
-    schema::{ChatCompletionChoice, Embedding},
+use spin_world::{
+    async_trait,
+    v2::llm::{self as wasi_llm},
 };
 
+use crate::schema::{ChatCompletionChoice, Embedding};
+
 mod default;
 mod open_ai;
 mod schema;
 
-#[derive(Clone)]
-pub enum Agent {
-    //TODO: Joshua: Naming??!
-    Default {
-        auth_token: String,
-        url: Url,
-        client: Option<Client>,
-    },
-    OpenAI {
-        auth_token: String,
-        url: Url,
-        client: Option<Client>,
-    },
+pub struct RemoteHttpLlmEngine {
+    worker: Box<dyn LlmWorker>,
 }
 
-impl Agent {
-    pub fn from(url: Url, auth_token: String, agent: Option<CustomLlm>) -> Self {
-        match agent {
-            Some(CustomLlm::OpenAi) => Agent::OpenAI {
-                auth_token,
-                url,
-                client: None,
-            },
-            _ => Agent::Default {
-                auth_token,
-                url,
-                client: None,
-            },
-        }
+impl RemoteHttpLlmEngine {
+    pub fn new(url: Url, auth_token: String, custom_llm: CustomLlm) -> Self {
+        let worker: Box<dyn LlmWorker> = match custom_llm {
+            CustomLlm::OpenAi => Box::new(open_ai::OpenAIAgentEngine::new(auth_token, url, None)),
+            CustomLlm::Default => Box::new(default::DefaultAgentEngine::new(auth_token, url, None)),
+        };
+        Self { worker }
     }
 }
 
-#[derive(Clone)]
-pub struct RemoteHttpLlmEngine {
-    agent: Agent,
+#[async_trait]
+pub trait LlmWorker: Send + Sync {
+    async fn infer(
+        &mut self,
+        model: wasi_llm::InferencingModel,
+        prompt: String,
+        params: wasi_llm::InferencingParams,
+    ) -> Result<wasi_llm::InferencingResult, wasi_llm::Error>;
+
+    async fn generate_embeddings(
+        &mut self,
+        model: wasi_llm::EmbeddingModel,
+        data: Vec<String>,
+    ) -> Result<wasi_llm::EmbeddingsResult, wasi_llm::Error>;
+
+    fn url(&self) -> Url;
 }
 
 #[derive(Serialize)]
@@ -150,74 +144,25 @@ struct OpenAIEmbeddingUsage {
 }
 
 impl RemoteHttpLlmEngine {
-    pub fn new(url: Url, auth_token: String, agent: Option<CustomLlm>) -> Self {
-        RemoteHttpLlmEngine {
-            agent: Agent::from(url, auth_token, agent),
-        }
-    }
-
     pub async fn infer(
         &mut self,
         model: wasi_llm::InferencingModel,
         prompt: String,
         params: wasi_llm::InferencingParams,
     ) -> Result<wasi_llm::InferencingResult, wasi_llm::Error> {
-        match &self.agent {
-            Agent::OpenAI {
-                auth_token,
-                url,
-                client,
-            } => {
-                OpenAIAgentEngine::infer(auth_token, url, client.clone(), model, prompt, params)
-                    .await
-            }
-            Agent::Default {
-                auth_token,
-                url,
-                client,
-            } => {
-                DefaultAgentEngine::infer(auth_token, url, client.clone(), model, prompt, params)
-                    .await
-            }
-        }
+        self.worker.infer(model, prompt, params).await
     }
 
     pub async fn generate_embeddings(
         &mut self,
         model: wasi_llm::EmbeddingModel,
         data: Vec<String>,
     ) -> Result<wasi_llm::EmbeddingsResult, wasi_llm::Error> {
-        match &self.agent {
-            Agent::OpenAI {
-                auth_token,
-                url,
-                client,
-            } => {
-                OpenAIAgentEngine::generate_embeddings(auth_token, url, client.clone(), model, data)
-                    .await
-            }
-            Agent::Default {
-                auth_token,
-                url,
-                client,
-            } => {
-                DefaultAgentEngine::generate_embeddings(
-                    auth_token,
-                    url,
-                    client.clone(),
-                    model,
-                    data,
-                )
-                .await
-            }
-        }
+        self.worker.generate_embeddings(model, data).await
     }
 
     pub fn url(&self) -> Url {
-        match &self.agent {
-            Agent::OpenAI { url, .. } => url.clone(),
-            Agent::Default { url, .. } => url.clone(),
-        }
+        self.worker.url()
     }
 }
 
@@ -267,19 +212,11 @@ impl From<CreateEmbeddingResponse> for wasi_llm::EmbeddingsResult {
     }
 }
 
-#[derive(Debug, serde::Deserialize, PartialEq)]
+#[derive(Debug, Default, serde::Deserialize, PartialEq)]
+#[serde(rename_all = "snake_case")]
 pub enum CustomLlm {
     /// Compatible with OpenAI's API alongside some other LLMs
     OpenAi,
-}
-
-impl TryFrom<&str> for CustomLlm {
-    type Error = anyhow::Error;
-
-    fn try_from(value: &str) -> Result<Self, Self::Error> {
-        match value.to_lowercase().as_str() {
-            "open_ai" | "openai" => Ok(CustomLlm::OpenAi),
-            _ => Err(anyhow::anyhow!("Invalid custom LLM: {}", value)),
-        }
-    }
+    #[default]
+    Default,
 }
diff --git a/crates/llm-remote-http/src/open_ai.rs b/crates/llm-remote-http/src/open_ai.rs
