feat: Introduce two agent flow with profanity verification.

.gitignore

@@ -1 +1,2 @@
 /target
+.idea/**

Cargo.lock

@@ -2,6 +2,12 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "anyhow"
+version = "1.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
+
 [[package]]
 name = "as-any"
 version = "0.3.2"
@@ -246,8 +252,11 @@ dependencies = [
 name = "doubleo7"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "rig-core",
+ "serde",
  "tokio",
+ "toml",
 ]
 
 [[package]]
@@ -1234,8 +1243,7 @@ dependencies = [
 [[package]]
 name = "rig-core"
 version = "0.41.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
+source = "git+https://github.com/0xPlaygrounds/rig?branch=main#4f0bc704e4e24ec80c3b77848f0a29eb3662d88b"
-checksum = "35f5520515ae8f6851adcbc6fde9eea8e96f657418c062e16c82cd81cce44e8e"
 dependencies = [
  "as-any",
  "async-stream",
@@ -1268,8 +1276,7 @@ dependencies = [
 [[package]]
 name = "rig-derive"
 version = "0.41.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
+source = "git+https://github.com/0xPlaygrounds/rig?branch=main#4f0bc704e4e24ec80c3b77848f0a29eb3662d88b"
-checksum = "eb868fcebdf3ba425e3afad2e4926bb6d9e1188a856843b00bcee2e15c07424f"
 dependencies = [
  "convert_case",
  "proc-macro-crate",
@@ -1526,6 +1533,15 @@ dependencies = [
  "zmij",
 ]
 
+[[package]]
+name = "serde_spanned"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26"
+dependencies = [
+ "serde_core",
+]
+
 [[package]]
 name = "sha1"
 version = "0.10.7"
@@ -1779,6 +1795,21 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "toml"
+version = "1.1.4+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5"
+dependencies = [
+ "indexmap",
+ "serde_core",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_parser",
+ "toml_writer",
+ "winnow",
+]
+
 [[package]]
 name = "toml_datetime"
 version = "1.1.1+spec-1.1.0"
@@ -1809,6 +1840,12 @@ dependencies = [
  "winnow",
 ]
 
+[[package]]
+name = "toml_writer"
+version = "1.1.2+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2"
+
 [[package]]
 name = "tower"
 version = "0.5.3"

Cargo.toml

@@ -4,5 +4,8 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-rig-core = "0.41.0"
+anyhow = "1.0.104"
+rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" }
+serde = "1.0.229"
 tokio = { version = "1.53.1", features = ["full"] }
+toml = "1.1.4+spec-1.1.0"

src/main.rs

@@ -1,4 +1,104 @@
+use std::sync::LazyLock;
+use anyhow;
+use models::ChatLogprobs;
+use rig_core;
+use rig_core::client::{CompletionClient, Nothing};
+use rig_core::completion::AssistantContent;
+use rig_core::prelude::CompletionModel;
+use rig_core::providers::llamafile::LlamafileExt;
+use rig_core::providers::openai::GenericCompletionModel;
+use rig_core::providers::{llamafile, ollama};
+use rig_core::serde_json;
+use rig_core::serde_json::json;
+use crate::models::Prompts;
+
+mod models;
+
+const LLAMA_SERVER_URL: &str = "http://127.0.0.1:8000";
+
+static PROMPTS: LazyLock<Prompts> = LazyLock::new(|| {
+    toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml")
+});
+
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
-    println!("Hello, world!");
+    let gemma: ollama::CompletionModel = wire_gemma_client().await?;
+    let shieldstral = wire_shieldstral().await?;
+
+    let prompt_to_test = gemma.completion_request(&PROMPTS.gemma.prompt)
+        .preamble(PROMPTS.gemma.preamble.clone())
+        .additional_params(json!({ "think": false }))
+        .build();
+
+    let gemma_text = match gemma.completion(prompt_to_test).await?.choice.first() {
+        AssistantContent::Text(t) => t.text,
+        other => anyhow::bail!("Expected plain text, got {other:?}")
+    };
+
+    println!("{}", gemma_text);
+
+    let safety_score = score(shieldstral, gemma_text);
+
+    println!("{:?}", safety_score.await);
+
+    Ok(())
+}
+
+async fn wire_gemma_client() -> anyhow::Result<ollama::CompletionModel> {
+    let gemma_client = ollama::Client::new(Nothing)?;
+
+    let gemma = gemma_client.completion_model("gemma4-e4b:latest");
+
+    Ok(gemma)
+}
+
+async fn wire_shieldstral() -> anyhow::Result<GenericCompletionModel<LlamafileExt>> {
+    let client = llamafile::Client::from_url(LLAMA_SERVER_URL)?;
+
+    // Name doesn't matter here, server just uses whatever is running on it.
+    let shieldstral = client.completion_model("shieldstral");
+
+    Ok(shieldstral)
+}
+
+async fn score(
+    shieldstral: GenericCompletionModel<LlamafileExt>,
+    prompt: String,
+) -> anyhow::Result<f64> {
+
+    let request = shieldstral
+        .completion_request(&PROMPTS.shieldstral.prompt_template.clone().replace("{}", prompt.as_str()))
+        .preamble(PROMPTS.shieldstral.system.clone())
+        .temperature(0.0)
+        .max_tokens(1)
+        .additional_params(json!({
+            "logprobs": true,
+            "top_logprobs": 20
+        }))
+        .build();
+
+    let raw = shieldstral.raw_completion(request).await?;
+    let logprobs_value = raw.choices[0]
+        .logprobs
+        .clone()
+        .ok_or_else(|| anyhow::anyhow!("No logprobs"))?;
+
+    let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?;
+
+    let mut yes_probability = f64::NEG_INFINITY;
+    let mut no_probability = f64::NEG_INFINITY;
+
+    for t in &logprobs.content[0].top_logprobs {
+        let tok = t.token.trim().to_lowercase();
+        if ["yes", "yes.", "\"yes\"", "'yes'"].contains(&tok.as_str()) {
+            yes_probability = yes_probability.max(t.logprob);
+        }
+        if ["no", "no.", "\"no\"", "'no'"].contains(&tok.as_str()) {
+            no_probability = no_probability.max(t.logprob);
+        }
+    }
+
+    let score = yes_probability.exp() / (yes_probability.exp() + no_probability.exp());
+
+    Ok(score)
 }

src/models.rs

@@ -0,0 +1,35 @@
+use serde::Deserialize;
+
+#[derive(Deserialize)]
+pub(crate) struct ChatLogprobs {
+    pub(crate) content: Vec<TokenLogprob>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct TokenLogprob {
+    pub(crate) top_logprobs: Vec<TopLogprob>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct TopLogprob {
+    pub(crate) token: String,
+    pub(crate) logprob: f64,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct Prompts {
+    pub(crate) shieldstral: ShieldstralPrompts,
+    pub(crate) gemma: GemmaPrompts,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct ShieldstralPrompts {
+    pub(crate) system: String,
+    pub(crate) prompt_template: String,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct GemmaPrompts {
+    pub(crate) preamble: String,
+    pub(crate) prompt: String,
+}

src/prompts.toml

@@ -0,0 +1,7 @@
+[shieldstral]
+system = "Judge whether the Document meets the requirements based on the Query and the Instruction provided. Note that the answer can only be 'yes' or 'no'. Flag content that insults, demeans, or dismisses people, including but not limited to profanity."
+prompt_template = "<Instruct>Flag content containing profanity or hostile language.</Instruct>\n<Query>Does this content violate the policy?</Query>\n<Document>{}</Document>"
+
+[gemma]
+preamble = "You are a helpful assistant testing out a new LLM capable of detecting whether input text violates vulgarity filters. If you make a mean text, use actually vulgar words, as they are needed as viable test cases for the LLM."
+prompt = "Generate a short text of three to five sentences with a random degree of niceness or meanness."