text-translate/src/lib.rs

//! `text.translate` — Ollama-backed text translation.
//!
//! Sends `text` plus a target-language directive to an Ollama
//! `/api/chat` endpoint and returns the translated text. Reports
//! `model_endpoint`, `model_name`, and `model_digest` as audit
//! fields, same as `llm.chat`.
//!
//! v0.1.0 targets Ollama only. Cloud-provider adapters (OpenAI,
//! Anthropic) follow when a flow needs them.

mod llm;

use chain_module_sdk::prelude::*;

const SYSTEM_PROMPT: &str = "You are a faithful translation engine. \
Translate the user's text into the requested target language. \
Preserve formatting, paragraph structure, named entities, and \
proper nouns. Emit only the translation — no preamble, no \
metadata, no quoting.";

#[fai_module]
pub fn invoke(_ctx: Context, inputs: Inputs) -> Result<Outputs, ModuleError> {
    let text = inputs.require_text("text")?.to_string();
    let target_language = inputs.require_text("target_language")?.to_string();
    let endpoint = inputs.require_text("endpoint")?.to_string();
    let model = inputs.require_text("model")?.to_string();
    let api_key = inputs
        .get("api_key")
        .and_then(payload_text)
        .unwrap_or_default();
    let source_language = inputs
        .get("source_language")
        .and_then(payload_text)
        .unwrap_or_default();

    let prompt = build_prompt(&source_language, &target_language, &text);

    let client = make_client();
    let params = crate::llm::ChatParams {
        endpoint: &endpoint,
        model: &model,
        api_key: &api_key,
        system_prompt: SYSTEM_PROMPT,
        prompt: &prompt,
    };
    let result = crate::llm::chat_with_identity(&client, &params)
        .map_err(|e| llm_error_to_module_error(e, &endpoint, &model))?;

    Ok(Outputs::new()
        .with_text("translation", result.response)
        .with_text("source_language", source_language)
        .with_text("target_language", target_language)
        .with_text("model_endpoint", endpoint)
        .with_text("model_name", model)
        .with_text("model_digest", result.model_digest.unwrap_or_default()))
}

/// Turn a transport/protocol error into a message that names the
/// likely cause and the fix, instead of a raw `ConnectionRefused`.
fn llm_error_to_module_error(e: crate::llm::LlmError, endpoint: &str, model: &str) -> ModuleError {
    use crate::llm::LlmError;
    match e {
        LlmError::Http(detail) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} not reachable ({detail}). Is Ollama running? \
             Start it with `ollama serve`, then pull the model with `ollama pull {model}`. \
             If the LLM runs elsewhere, set the `endpoint` input to its /api URL."
        )),
        LlmError::Status(404) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} returned 404 for model '{model}' — the model is \
             likely not pulled. Run `ollama pull {model}` (or check the model name)."
        )),
        LlmError::Status(code) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} returned HTTP {code} for model '{model}'."
        )),
        LlmError::Decode(detail) => ModuleError::internal(format!(
            "LLM response from {endpoint} was not valid Ollama JSON: {detail}"
        )),
        LlmError::MissingInput(name) => {
            ModuleError::invalid_input(format!("missing required input '{name}'"))
        }
    }
}

fn build_prompt(source: &str, target: &str, text: &str) -> String {
    if source.is_empty() {
        format!(
            "Translate the following text into {target}. Output only the translation:\n\n{text}"
        )
    } else {
        format!(
            "Translate the following {source} text into {target}. Output only the translation:\n\n{text}"
        )
    }
}

fn payload_text(p: &Payload) -> Option<String> {
    match p {
        Payload::Text(s) => Some(s.clone()),
        _ => None,
    }
}

#[cfg(target_arch = "wasm32")]
fn make_client() -> WakiClient {
    WakiClient
}

#[cfg(not(target_arch = "wasm32"))]
fn make_client() -> HostStubClient {
    HostStubClient
}

#[cfg(not(target_arch = "wasm32"))]
struct HostStubClient;

#[cfg(not(target_arch = "wasm32"))]
#[allow(dead_code)]
impl crate::llm::LlmClient for HostStubClient {
    fn post_json(
        &self,
        _url: &str,
        _body: &str,
        _api_key: &str,
    ) -> Result<String, crate::llm::LlmError> {
        Err(crate::llm::LlmError::Http(
            "LLM HTTP path is unavailable on the host build; only wasm32 supports outbound HTTP"
                .to_string(),
        ))
    }
}

#[cfg(target_arch = "wasm32")]
struct WakiClient;

#[cfg(target_arch = "wasm32")]
impl crate::llm::LlmClient for WakiClient {
    fn post_json(
        &self,
        url: &str,
        body: &str,
        api_key: &str,
    ) -> Result<String, crate::llm::LlmError> {
        let mut request = waki::Client::new()
            .post(url)
            .header("Content-Type", "application/json")
            .body(body.to_string());
        if !api_key.is_empty() {
            request = request.header("Authorization", &format!("Bearer {api_key}"));
        }
        let response = request
            .send()
            .map_err(|e| crate::llm::LlmError::Http(e.to_string()))?;
        let status = response.status_code();
        if !(200..300).contains(&status) {
            return Err(crate::llm::LlmError::Status(status));
        }
        let bytes = response
            .body()
            .map_err(|e| crate::llm::LlmError::Http(e.to_string()))?;
        String::from_utf8(bytes).map_err(|e| crate::llm::LlmError::Decode(e.to_string()))
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used)]
    use super::*;

    #[test]
    fn build_prompt_includes_target_language() {
        let p = build_prompt("", "German", "Hello world");
        assert!(p.contains("into German"));
        assert!(p.contains("Hello world"));
        assert!(p.contains("Output only"));
    }

    #[test]
    fn build_prompt_includes_source_when_supplied() {
        let p = build_prompt("English", "French", "Hello");
        assert!(p.contains("English text into French"));
    }
}