wow

src/main.rs

@@ -0,0 +1,297 @@
+use std::{
+    fs,
+    io::{BufRead, BufReader, ErrorKind, Read, Write},
+    net::{TcpListener, TcpStream},
+    path::Path,
+    process::Command,
+    sync::Arc,
+    thread,
+    time::Duration,
+};
+
+use serde::Deserialize;
+use serde_json::{json, Value};
+
+const PORT: u16 = 1108;
+const API: &str = "https://akiyama.mizuki.guru/api/post/random/url";
+const CONFIG_PATH: &str = "config.toml";
+const MODEL_ID: &str = "gemma-4-26b-a4b-it";
+const GENERATE_CONTENT_API: &str = "streamGenerateContent";
+
+#[derive(Clone, Deserialize)]
+struct Config {
+    google: GoogleConfig,
+}
+
+#[derive(Clone, Deserialize)]
+struct GoogleConfig {
+    #[serde(rename = "api-key")]
+    api_key: String,
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let config = Arc::new(load_config()?);
+    let listener = TcpListener::bind(("0.0.0.0", PORT))?;
+    println!("listening on :{PORT}");
+
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                let config = Arc::clone(&config);
+                thread::spawn(move || {
+                    if let Err(e) = handle_client(stream, &config) {
+                        eprintln!("client error: {e}");
+                    }
+                });
+            }
+            Err(e) => eprintln!("accept error: {e}"),
+        }
+    }
+
+    Ok(())
+}
+
+fn load_config() -> Result<Config, Box<dyn std::error::Error>> {
+    let config = fs::read_to_string(CONFIG_PATH)?;
+    let config = toml::from_str(&config)?;
+    Ok(config)
+}
+
+fn handle_client(
+    mut stream: TcpStream,
+    config: &Config,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let mut reader = BufReader::new(stream.try_clone()?);
+    let mut history = Vec::new();
+    let url = ureq::get(API).call()?.into_string()?;
+    let url = url.trim();
+
+    writeln!(
+        stream,
+        r#"███╗   ███╗██╗███████╗██╗   ██╗██╗  ██╗██╗    ██████╗ ██╗   ██╗██████╗ ██╗   ██╗
+████╗ ████║██║╚══███╔╝██║   ██║██║ ██╔╝██║   ██╔════╝ ██║   ██║██╔══██╗██║   ██║
+██╔████╔██║██║  ███╔╝ ██║   ██║█████╔╝ ██║   ██║  ███╗██║   ██║██████╔╝██║   ██║
+██║╚██╔╝██║██║ ███╔╝  ██║   ██║██╔═██╗ ██║   ██║   ██║██║   ██║██╔══██╗██║   ██║
+██║ ╚═╝ ██║██║███████╗╚██████╔╝██║  ██╗██║██╗╚██████╔╝╚██████╔╝██║  ██║╚██████╔╝
+╚═╝     ╚═╝╚═╝╚══════╝ ╚═════╝ ╚═╝  ╚═╝╚═╝╚═╝ ╚═════╝  ╚═════╝ ╚═╝  ╚═╝ ╚═════╝
+"#,
+    )?;
+
+    writeln!(stream, "URL: {url}")?;
+    writeln!(stream)?;
+
+    let mut response = ureq::get(url).call()?.into_reader();
+
+    let mut file = tempfile::Builder::new().suffix(".img").tempfile()?;
+
+    let mut buf = Vec::new();
+    response.read_to_end(&mut buf)?;
+    file.write_all(&buf)?;
+
+    write_image(&mut stream, file.path())?;
+
+    loop {
+        writeln!(stream)?;
+        write!(stream, "> ")?;
+        stream.flush()?;
+
+        let Some(prompt) = read_prompt(&mut reader, Duration::from_secs(300))? else {
+            break;
+        };
+
+        if matches!(prompt.as_str(), "exit" | "quit" | "종료") {
+            writeln!(stream, "bye~")?;
+            break;
+        }
+
+        writeln!(stream)?;
+        write!(stream, ": ")?;
+        stream.flush()?;
+
+        match generate_gemma_content(config, &mut history, &prompt) {
+            Ok(text) => writeln!(stream, "{text}")?,
+            Err(e) => writeln!(stream, "Gemma request failed: {e}")?,
+        }
+    }
+
+    Ok(())
+}
+
+fn write_image(stream: &mut TcpStream, image_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
+    let output = Command::new("chafa")
+        .args([
+            "-f",
+            "symbols",
+            "--colors=256",
+            "--size=100x60",
+            image_path.to_str().unwrap(),
+        ])
+        .output();
+
+    match output {
+        Ok(output) => {
+            stream.write_all(&output.stdout)?;
+
+            if !output.status.success() {
+                stream.write_all(&output.stderr)?;
+            }
+        }
+        Err(e) if e.kind() == ErrorKind::NotFound => {
+            writeln!(stream, "chafa not found; cannot render image in terminal.")?;
+        }
+        Err(e) => return Err(e.into()),
+    }
+
+    Ok(())
+}
+
+fn read_prompt(
+    reader: &mut BufReader<TcpStream>,
+    timeout: Duration,
+) -> Result<Option<String>, Box<dyn std::error::Error>> {
+    reader.get_ref().set_read_timeout(Some(timeout))?;
+
+    let mut line = String::new();
+    let read_result = reader.read_line(&mut line);
+
+    reader.get_ref().set_read_timeout(None)?;
+
+    match read_result {
+        Ok(0) => Ok(None),
+        Ok(_) => {
+            let prompt = line.trim().to_string();
+            Ok((!prompt.is_empty()).then_some(prompt))
+        }
+        Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => Ok(None),
+        Err(e) => Err(e.into()),
+    }
+}
+
+fn generate_gemma_content(
+    config: &Config,
+    history: &mut Vec<Value>,
+    prompt: &str,
+) -> Result<String, Box<dyn std::error::Error>> {
+    let mut contents = history.clone();
+    let user_content = content("user", prompt);
+    contents.push(user_content.clone());
+
+    let request = json!({
+        "contents": contents,
+        "generationConfig": {
+            "thinkingConfig": {
+                "thinkingLevel": "MINIMAL"
+            }
+        },
+        "safetySettings": [
+            {
+                "category": "HARM_CATEGORY_HARASSMENT",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_HATE_SPEECH",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                "threshold": "BLOCK_NONE"
+            }
+        ],
+        "tools": [
+            {
+                "googleSearch": {}
+            }
+        ],
+        "systemInstruction": {
+            "parts": [
+                {
+                    "text": "당신은 프로젝트 세카이의 아키야마 미즈키입니다."
+                }
+            ]
+        }
+    });
+
+    let url = format!(
+        "https://generativelanguage.googleapis.com/v1beta/models/{MODEL_ID}:{GENERATE_CONTENT_API}?key={}",
+        config.google.api_key
+    );
+
+    let response = match ureq::post(&url)
+        .set("Content-Type", "application/json")
+        .send_string(&request.to_string())
+    {
+        Ok(response) => response,
+        Err(ureq::Error::Status(code, response)) => {
+            let body = response
+                .into_string()
+                .unwrap_or_else(|_| "failed to read error body".to_string());
+            return Err(std::io::Error::new(
+                ErrorKind::Other,
+                format!("Gemma API returned HTTP {code}: {body}"),
+            )
+            .into());
+        }
+        Err(e) => return Err(e.into()),
+    };
+
+    let body = response.into_string()?;
+    let text = extract_gemma_text(&body).unwrap_or(body);
+
+    history.push(user_content);
+    history.push(content("model", &text));
+
+    Ok(text)
+}
+
+fn content(role: &str, text: &str) -> Value {
+    json!({
+        "role": role,
+        "parts": [
+            {
+                "text": text
+            }
+        ]
+    })
+}
+
+fn extract_gemma_text(body: &str) -> Option<String> {
+    let value: Value = serde_json::from_str(body).ok()?;
+    let mut texts = Vec::new();
+
+    match value {
+        Value::Array(chunks) => {
+            for chunk in &chunks {
+                collect_candidate_text(chunk, &mut texts);
+            }
+        }
+        value => collect_candidate_text(&value, &mut texts),
+    }
+
+    (!texts.is_empty()).then(|| texts.join(""))
+}
+
+fn collect_candidate_text(value: &Value, texts: &mut Vec<String>) {
+    let Some(candidates) = value.get("candidates").and_then(Value::as_array) else {
+        return;
+    };
+
+    for candidate in candidates {
+        let Some(parts) = candidate
+            .get("content")
+            .and_then(|content| content.get("parts"))
+            .and_then(Value::as_array)
+        else {
+            continue;
+        };
+
+        for part in parts {
+            if let Some(text) = part.get("text").and_then(Value::as_str) {
+                texts.push(text.to_string());
+            }
+        }
+    }
+}