text-translate/src/lib.rs

//! `text.translate` — LLM-backed text translation.
//!
//! Sends `text` plus a target-language directive to an LLM chat
//! endpoint and returns the translated text. Reports
//! `model_endpoint`, `model_name`, and `model_digest` as audit
//! fields, same as `llm.chat`.
//!
//! The wire format is selected by the optional `api` input:
//! `ollama` (default, unchanged v0.1.x behavior), `openai`
//! (OpenAI-compatible `/v1/chat/completions` — vLLM, LM Studio,
//! LiteLLM, cloud OpenAI), or `anthropic` (Messages API). The
//! client logic in `llm.rs` is kept in lockstep with `llm.chat`.

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.";

#[chain_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 api_raw = inputs.get("api").and_then(payload_text).unwrap_or_default();
    let api =
        crate::llm::Api::parse(&api_raw).map_err(|e| ModuleError::invalid_input(e.to_string()))?;

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

    let client = make_client();
    let params = crate::llm::ChatParams {
        api,
        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, api, &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`.
/// The hints are api-specific: an Ollama connect failure almost
/// always means Ollama isn't running or the model isn't pulled,
/// while cloud/vLLM failures are usually endpoint or key issues.
/// The message must never contain the api_key.
fn llm_error_to_module_error(
    e: crate::llm::LlmError,
    api: crate::llm::Api,
    endpoint: &str,
    model: &str,
) -> ModuleError {
    use crate::llm::{Api, LlmError};
    match (api, e) {
        (Api::Ollama, 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."
        )),
        (Api::Openai, LlmError::Http(detail)) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} not reachable ({detail}). Expected an \
             OpenAI-compatible server (OpenAI, vLLM, ...) at a /v1/chat/completions URL."
        )),
        (Api::Anthropic, LlmError::Http(detail)) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} not reachable ({detail}). Expected the \
             Anthropic Messages API at a /v1/messages URL."
        )),
        (Api::Ollama, 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(401)) | (_, LlmError::Status(403)) => ModuleError::internal(format!(
            "LLM endpoint {endpoint} rejected the request as unauthorized — check the \
             `api_key` input."
        )),
        (_, 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} did not match the expected schema: {detail}"
        )),
        (_, LlmError::MissingInput(name)) => {
            ModuleError::invalid_input(format!("missing required input '{name}'"))
        }
        (_, LlmError::UnsupportedApi(raw)) => ModuleError::invalid_input(format!(
            "unsupported api '{raw}' (expected: ollama, openai, anthropic)"
        )),
    }
}

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,
        _headers: &[(&'static str, String)],
    ) -> 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,
        headers: &[(&'static str, String)],
    ) -> Result<String, crate::llm::LlmError> {
        let mut request = waki::Client::new()
            .post(url)
            .header("Content-Type", "application/json")
            .body(body.to_string());
        for (name, value) in headers {
            request = request.header(*name, value);
        }
        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"));
    }

    #[test]
    fn empty_api_input_defaults_to_ollama() {
        assert_eq!(crate::llm::Api::parse("").unwrap(), crate::llm::Api::Ollama);
    }

    #[test]
    fn openai_api_is_accepted_for_vllm_endpoints() {
        assert_eq!(
            crate::llm::Api::parse("openai").unwrap(),
            crate::llm::Api::Openai
        );
    }
}