text.anonymize.legal-de 0.1.0: WASM bridge to the judge-ner host service

NER anonymization of German legal texts (plain text + DOCX) via the
JuraNER model by Harshil Darji (MIT). The transformer model runs in
the judge-ner host service; this module forwards documents over
loopback HTTP and maps responses onto the capability contract.

src/anonymize.rs

@@ -0,0 +1,439 @@
+//! Request/response mapping between the frozen
+//! `text.anonymize.legal-de` contract and the judge-ner host
+//! service HTTP API.
+//!
+//! All HTTP I/O lives behind the `HttpClient` trait so unit tests
+//! exercise request building, response parsing and error mapping on
+//! the host without a network.
+
+use base64::Engine as _;
+use base64::engine::general_purpose::STANDARD as B64;
+use serde::Deserialize;
+use serde_json::json;
+
+/// Base URL used when the optional `endpoint` input is absent/empty.
+pub const DEFAULT_ENDPOINT: &str = "http://127.0.0.1:8756";
+
+pub const DOCX_MIME: &str =
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
+pub const TEXT_MIME: &str = "text/plain; charset=utf-8";
+
+#[derive(Debug, thiserror::Error)]
+pub enum AnonError {
+    /// Caller-side problem (bad request JSON, non-UTF-8 text
+    /// document, service rejected the input as invalid).
+    #[error("invalid input: {0}")]
+    InvalidInput(String),
+    /// Transport failure or unexpected service behaviour.
+    #[error("{0}")]
+    Service(String),
+}
+
+/// Minimal HTTP response the mapping layer needs.
+pub struct HttpResponse {
+    pub status: u16,
+    pub body: String,
+}
+
+pub trait HttpClient {
+    fn post_json(&self, url: &str, body: &str) -> Result<HttpResponse, AnonError>;
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Mode {
+    Text,
+    Docx,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct Request {
+    pub mode: Mode,
+    #[serde(default)]
+    pub options: serde_json::Value,
+}
+
+/// Parse the `request` input JSON.
+pub fn parse_request(raw: &str) -> Result<Request, AnonError> {
+    let mut req: Request = serde_json::from_str(raw).map_err(|e| {
+        AnonError::InvalidInput(format!(
+            "request must be {{\"mode\":\"text\"|\"docx\",\"options\":{{..}}}}: {e}"
+        ))
+    })?;
+    if req.options.is_null() {
+        req.options = json!({});
+    } else if !req.options.is_object() {
+        return Err(AnonError::InvalidInput(
+            "request.options must be a JSON object".to_string(),
+        ));
+    }
+    Ok(req)
+}
+
+/// Normalize the endpoint override: empty -> default, strip a
+/// trailing slash so path joining stays predictable.
+pub fn resolve_endpoint(endpoint_input: &str) -> String {
+    let trimmed = endpoint_input.trim();
+    let base = if trimmed.is_empty() {
+        DEFAULT_ENDPOINT
+    } else {
+        trimmed
+    };
+    base.trim_end_matches('/').to_string()
+}
+
+/// Outcome of one anonymization call, ready to be lifted into
+/// module `Outputs` by the caller.
+#[derive(Debug)]
+pub struct AnonOutcome {
+    /// JSON string for the `result` output.
+    pub result_json: String,
+    /// Bytes for the `document` output.
+    pub document: Vec<u8>,
+    /// MIME type of `document`.
+    pub document_mime: &'static str,
+}
+
+pub fn run<C: HttpClient>(
+    client: &C,
+    request: &Request,
+    document: &[u8],
+    endpoint_input: &str,
+) -> Result<AnonOutcome, AnonError> {
+    let base = resolve_endpoint(endpoint_input);
+    match request.mode {
+        Mode::Text => run_text(client, &base, &request.options, document),
+        Mode::Docx => run_docx(client, &base, &request.options, document),
+    }
+}
+
+fn run_text<C: HttpClient>(
+    client: &C,
+    base: &str,
+    options: &serde_json::Value,
+    document: &[u8],
+) -> Result<AnonOutcome, AnonError> {
+    let text = std::str::from_utf8(document).map_err(|_| {
+        AnonError::InvalidInput(
+            "document is not valid UTF-8 text (mode=text expects a UTF-8 text payload)"
+                .to_string(),
+        )
+    })?;
+    let url = format!("{base}/anonymize/text");
+    let body = serde_json::to_string(&json!({ "text": text, "options": options }))
+        .map_err(|e| AnonError::Service(format!("serialize request: {e}")))?;
+    let response = post_expect_json(client, &url, &body)?;
+
+    let anonymized = response
+        .get("text")
+        .and_then(|t| t.as_str())
+        .ok_or_else(|| AnonError::Service(
+            "judge-ner response is missing the 'text' field".to_string(),
+        ))?
+        .to_string();
+    let result_json = build_result(&response, Some(&anonymized))?;
+
+    Ok(AnonOutcome {
+        result_json,
+        document: anonymized.into_bytes(),
+        document_mime: TEXT_MIME,
+    })
+}
+
+fn run_docx<C: HttpClient>(
+    client: &C,
+    base: &str,
+    options: &serde_json::Value,
+    document: &[u8],
+) -> Result<AnonOutcome, AnonError> {
+    let url = format!("{base}/anonymize/docx");
+    let body = serde_json::to_string(&json!({
+        "docx_base64": B64.encode(document),
+        "options": options,
+    }))
+    .map_err(|e| AnonError::Service(format!("serialize request: {e}")))?;
+    let response = post_expect_json(client, &url, &body)?;
+
+    let docx_b64 = response
+        .get("docx_base64")
+        .and_then(|d| d.as_str())
+        .ok_or_else(|| AnonError::Service(
+            "judge-ner response is missing the 'docx_base64' field".to_string(),
+        ))?;
+    let docx_bytes = B64.decode(docx_b64).map_err(|e| {
+        AnonError::Service(format!("judge-ner returned invalid base64 DOCX: {e}"))
+    })?;
+    let result_json = build_result(&response, None)?;
+
+    Ok(AnonOutcome {
+        result_json,
+        document: docx_bytes,
+        document_mime: DOCX_MIME,
+    })
+}
+
+/// Build the `result` output JSON from the service response.
+///
+/// Frozen contract: `{"replacements": [...]}` plus `"text"` only in
+/// text mode. `statistics` / `rubrum` are passed through as optional
+/// extras when the service supplied them.
+fn build_result(
+    response: &serde_json::Value,
+    text: Option<&str>,
+) -> Result<String, AnonError> {
+    let replacements = response
+        .get("replacements")
+        .cloned()
+        .ok_or_else(|| AnonError::Service(
+            "judge-ner response is missing the 'replacements' field".to_string(),
+        ))?;
+    if !replacements.is_array() {
+        return Err(AnonError::Service(
+            "judge-ner 'replacements' field is not an array".to_string(),
+        ));
+    }
+
+    let mut result = serde_json::Map::new();
+    result.insert("replacements".to_string(), replacements);
+    if let Some(t) = text {
+        result.insert("text".to_string(), json!(t));
+    }
+    for passthrough in ["statistics", "rubrum"] {
+        if let Some(v) = response.get(passthrough) {
+            result.insert(passthrough.to_string(), v.clone());
+        }
+    }
+    serde_json::to_string(&serde_json::Value::Object(result))
+        .map_err(|e| AnonError::Service(format!("serialize result: {e}")))
+}
+
+/// POST and decode, translating HTTP-level failures into contract
+/// errors. FastAPI signals caller errors as 4xx with a JSON body
+/// `{"detail": "..."}` — surface those as InvalidInput so the flow
+/// engine reports them as the caller's problem, not the module's.
+fn post_expect_json<C: HttpClient>(
+    client: &C,
+    url: &str,
+    body: &str,
+) -> Result<serde_json::Value, AnonError> {
+    let response = client.post_json(url, body)?;
+    if !(200..300).contains(&response.status) {
+        let detail = serde_json::from_str::<serde_json::Value>(&response.body)
+            .ok()
+            .and_then(|v| v.get("detail").and_then(|d| d.as_str().map(String::from)))
+            .unwrap_or_else(|| truncate(&response.body, 300));
+        let message = format!(
+            "judge-ner at {url} answered HTTP {status}: {detail}",
+            status = response.status,
+        );
+        return if (400..500).contains(&response.status) && response.status != 404 {
+            Err(AnonError::InvalidInput(message))
+        } else {
+            Err(AnonError::Service(message))
+        };
+    }
+    serde_json::from_str(&response.body).map_err(|e| {
+        AnonError::Service(format!("judge-ner returned non-JSON body from {url}: {e}"))
+    })
+}
+
+fn truncate(s: &str, max: usize) -> String {
+    if s.len() <= max {
+        s.to_string()
+    } else {
+        let mut cut = max;
+        while !s.is_char_boundary(cut) {
+            cut -= 1;
+        }
+        format!("{}…", &s[..cut])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
+    use super::*;
+    use std::cell::RefCell;
+
+    struct MockClient {
+        responses: RefCell<Vec<Result<HttpResponse, AnonError>>>,
+        seen: RefCell<Vec<(String, String)>>,
+    }
+
+    impl MockClient {
+        fn new(responses: Vec<Result<HttpResponse, AnonError>>) -> Self {
+            Self {
+                responses: RefCell::new(responses),
+                seen: RefCell::new(Vec::new()),
+            }
+        }
+    }
+
+    impl HttpClient for MockClient {
+        fn post_json(&self, url: &str, body: &str) -> Result<HttpResponse, AnonError> {
+            self.seen
+                .borrow_mut()
+                .push((url.to_string(), body.to_string()));
+            self.responses
+                .borrow_mut()
+                .pop()
+                .unwrap_or(Err(AnonError::Service("no more mock responses".into())))
+        }
+    }
+
+    fn ok(body: &str) -> Result<HttpResponse, AnonError> {
+        Ok(HttpResponse {
+            status: 200,
+            body: body.to_string(),
+        })
+    }
+
+    #[test]
+    fn parse_request_accepts_text_mode_without_options() {
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        assert_eq!(req.mode, Mode::Text);
+        assert!(req.options.as_object().unwrap().is_empty());
+    }
+
+    #[test]
+    fn parse_request_accepts_docx_mode_with_options() {
+        let req =
+            parse_request(r#"{"mode":"docx","options":{"remove_rubrum":true}}"#).unwrap();
+        assert_eq!(req.mode, Mode::Docx);
+        assert_eq!(req.options["remove_rubrum"], true);
+    }
+
+    #[test]
+    fn parse_request_rejects_unknown_mode() {
+        assert!(matches!(
+            parse_request(r#"{"mode":"pdf"}"#),
+            Err(AnonError::InvalidInput(_))
+        ));
+    }
+
+    #[test]
+    fn parse_request_rejects_non_object_options() {
+        assert!(matches!(
+            parse_request(r#"{"mode":"text","options":[1,2]}"#),
+            Err(AnonError::InvalidInput(_))
+        ));
+    }
+
+    #[test]
+    fn resolve_endpoint_defaults_and_strips_slash() {
+        assert_eq!(resolve_endpoint(""), DEFAULT_ENDPOINT);
+        assert_eq!(resolve_endpoint("  "), DEFAULT_ENDPOINT);
+        assert_eq!(
+            resolve_endpoint("http://10.0.0.5:9000/"),
+            "http://10.0.0.5:9000"
+        );
+    }
+
+    #[test]
+    fn text_mode_maps_response_onto_contract() {
+        let service_response = r#"{
+            "text": "Herr ⟦PER1⟧.",
+            "replacements": [{"text":"Müller","label":"PER","start":5,"end":11,
+                              "anonymized":"⟦PER1⟧"}],
+            "statistics": {"total_entities": 1},
+            "processing_time": 0.4
+        }"#;
+        let client = MockClient::new(vec![ok(service_response)]);
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        let outcome = run(&client, &req, "Herr Müller.".as_bytes(), "").unwrap();
+
+        // URL + body
+        let (url, body) = client.seen.borrow()[0].clone();
+        assert_eq!(url, "http://127.0.0.1:8756/anonymize/text");
+        let sent: serde_json::Value = serde_json::from_str(&body).unwrap();
+        assert_eq!(sent["text"], "Herr Müller.");
+        assert!(sent["options"].is_object());
+
+        // result output
+        let result: serde_json::Value = serde_json::from_str(&outcome.result_json).unwrap();
+        assert_eq!(result["text"], "Herr ⟦PER1⟧.");
+        assert_eq!(result["replacements"][0]["anonymized"], "⟦PER1⟧");
+        assert_eq!(result["statistics"]["total_entities"], 1);
+        assert!(result.get("rubrum").is_none());
+
+        // document output = anonymized UTF-8 text
+        assert_eq!(outcome.document, "Herr ⟦PER1⟧.".as_bytes());
+        assert_eq!(outcome.document_mime, TEXT_MIME);
+    }
+
+    #[test]
+    fn text_mode_rejects_non_utf8_document() {
+        let client = MockClient::new(vec![]);
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        let err = run(&client, &req, &[0xff, 0xfe, 0x00], "").unwrap_err();
+        assert!(matches!(err, AnonError::InvalidInput(_)));
+    }
+
+    #[test]
+    fn docx_mode_encodes_and_decodes_base64() {
+        let input_docx = b"PK\x03\x04fake-docx".to_vec();
+        let output_docx = b"PK\x03\x04anonymized".to_vec();
+        let service_response = serde_json::to_string(&serde_json::json!({
+            "docx_base64": B64.encode(&output_docx),
+            "replacements": [],
+            "statistics": {"total_entities": 0}
+        }))
+        .unwrap();
+        let client = MockClient::new(vec![ok(&service_response)]);
+        let req = parse_request(r#"{"mode":"docx","options":{"remove_rubrum":true}}"#)
+            .unwrap();
+        let outcome = run(&client, &req, &input_docx, "http://127.0.0.1:9999/").unwrap();
+
+        let (url, body) = client.seen.borrow()[0].clone();
+        assert_eq!(url, "http://127.0.0.1:9999/anonymize/docx");
+        let sent: serde_json::Value = serde_json::from_str(&body).unwrap();
+        assert_eq!(sent["docx_base64"], B64.encode(&input_docx));
+        assert_eq!(sent["options"]["remove_rubrum"], true);
+
+        let result: serde_json::Value = serde_json::from_str(&outcome.result_json).unwrap();
+        assert!(result.get("text").is_none(), "no 'text' key in docx mode");
+        assert_eq!(result["replacements"].as_array().unwrap().len(), 0);
+
+        assert_eq!(outcome.document, output_docx);
+        assert_eq!(outcome.document_mime, DOCX_MIME);
+    }
+
+    #[test]
+    fn service_4xx_maps_to_invalid_input_with_detail() {
+        let client = MockClient::new(vec![Ok(HttpResponse {
+            status: 422,
+            body: r#"{"detail":"text must not be empty"}"#.to_string(),
+        })]);
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        let err = run(&client, &req, b"", "").unwrap_err();
+        match err {
+            AnonError::InvalidInput(msg) => {
+                assert!(msg.contains("text must not be empty"), "msg = {msg}");
+                assert!(msg.contains("422"), "msg = {msg}");
+            }
+            other => panic!("expected InvalidInput, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn service_5xx_maps_to_service_error() {
+        let client = MockClient::new(vec![Ok(HttpResponse {
+            status: 503,
+            body: r#"{"detail":"NER model not available: ..."}"#.to_string(),
+        })]);
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        let err = run(&client, &req, b"hello", "").unwrap_err();
+        assert!(matches!(err, AnonError::Service(_)));
+    }
+
+    #[test]
+    fn missing_replacements_is_a_service_error() {
+        let client = MockClient::new(vec![ok(r#"{"text":"x"}"#)]);
+        let req = parse_request(r#"{"mode":"text"}"#).unwrap();
+        let err = run(&client, &req, b"x", "").unwrap_err();
+        match err {
+            AnonError::Service(msg) => assert!(msg.contains("replacements")),
+            other => panic!("expected Service, got {other:?}"),
+        }
+    }
+}

src/lib.rs

@@ -0,0 +1,124 @@
+//! `text.anonymize.legal-de` — NER anonymization of German legal
+//! texts, bridged to the `judge-ner` host service.
+//!
+//! The heavy lifting (JuraNER transformers model, DOCX rewriting)
+//! happens in the judge-ner container — see
+//! `fai_judge/host_services/judge-ner/`. This module forwards the
+//! document over loopback HTTP and maps the response onto the frozen
+//! capability contract:
+//!
+//! inputs:  `request` (json)  {"mode":"text"|"docx","options":{...}}
+//!          `document` (bytes) UTF-8 text or DOCX
+//!          `endpoint` (text, optional) service base URL override
+//! outputs: `result` (json)  {"replacements":[...], "text" only for
+//!          mode=text}
+//!          `document` (bytes) anonymized text/DOCX
+
+mod anonymize;
+
+pub use anonymize::DEFAULT_ENDPOINT;
+
+use chain_module_sdk::prelude::*;
+
+use crate::anonymize::{AnonError, HttpClient, HttpResponse};
+
+#[fai_module]
+pub fn invoke(_ctx: Context, inputs: Inputs) -> Result<Outputs, ModuleError> {
+    let request_raw = read_json_or_text(&inputs, "request")?;
+    let request = anonymize::parse_request(&request_raw).map_err(to_module_error)?;
+    let document = inputs.require_bytes("document")?;
+    let endpoint = inputs
+        .get("endpoint")
+        .and_then(payload_text)
+        .unwrap_or_default();
+
+    let client = make_client();
+    let outcome = anonymize::run(&client, &request, &document.data, &endpoint)
+        .map_err(to_module_error)?;
+
+    Ok(Outputs::new()
+        .with_json_str("result", outcome.result_json)
+        .with_bytes("document", outcome.document_mime, outcome.document))
+}
+
+fn to_module_error(err: AnonError) -> ModuleError {
+    match err {
+        AnonError::InvalidInput(msg) => ModuleError::invalid_input(msg),
+        AnonError::Service(msg) => ModuleError::internal(msg),
+    }
+}
+
+fn payload_text(p: &Payload) -> Option<String> {
+    match p {
+        Payload::Text(s) => Some(s.clone()),
+        _ => None,
+    }
+}
+
+/// Accept the request payload either as `json` (the documented form)
+/// or `text` (the form gRPC clients without first-class JSON-payload
+/// support fall through to). Both carry the same JSON string.
+fn read_json_or_text(inputs: &Inputs, name: &str) -> Result<String, ModuleError> {
+    if let Ok(s) = inputs.require_json_str(name) {
+        return Ok(s.to_string());
+    }
+    if let Ok(s) = inputs.require_text(name) {
+        return Ok(s.to_string());
+    }
+    Err(ModuleError::invalid_input(format!(
+        "input '{name}' must be json or text (carrying a JSON string)"
+    )))
+}
+
+#[cfg(target_arch = "wasm32")]
+fn make_client() -> WakiClient {
+    WakiClient
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+fn make_client() -> HostStubClient {
+    HostStubClient
+}
+
+/// Host builds (unit tests, `cargo test`) have no outbound HTTP —
+/// only the wasm32 build talks to the service.
+#[cfg(not(target_arch = "wasm32"))]
+struct HostStubClient;
+
+#[cfg(not(target_arch = "wasm32"))]
+#[allow(dead_code)]
+impl HttpClient for HostStubClient {
+    fn post_json(&self, _url: &str, _body: &str) -> Result<HttpResponse, AnonError> {
+        Err(AnonError::Service(
+            "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 HttpClient for WakiClient {
+    fn post_json(&self, url: &str, body: &str) -> Result<HttpResponse, AnonError> {
+        let response = waki::Client::new()
+            .post(url)
+            .header("Content-Type", "application/json")
+            .body(body.to_string())
+            .send()
+            .map_err(|e| {
+                AnonError::Service(format!(
+                    "judge-ner service unreachable at {url}: {e} \
+                     (is the judge-ner container running?)"
+                ))
+            })?;
+        let status = response.status_code();
+        let bytes = response
+            .body()
+            .map_err(|e| AnonError::Service(format!("read response body: {e}")))?;
+        let body = String::from_utf8(bytes)
+            .map_err(|e| AnonError::Service(format!("response body is not UTF-8: {e}")))?;
+        Ok(HttpResponse { status, body })
+    }
+}