mini_ollama_client 0.1.0

//! # Ollama Client Library
//!
//! This library provides a simple client to interact with the Ollama server.
//! It allows sending requests to the server and receiving responses.

use std::io::{Read, Write};
use std::net::TcpStream;
use std::error::Error;

/// Sends a request to the Ollama server.
///
/// # Arguments
///
/// * `server` - A string slice that holds the server address.
/// * `prompt` - A string slice that holds the prompt to send to the server.
/// * `model` - An optional string slice that holds the model to use. If `None`, "phi3" will be used by default.
///
/// # Returns
///
/// A `Result` containing the response from the server as a `String` on success, or an error.
///
/// # Examples
///
/// ```
/// use ollama_client::send_request;
///
/// let server = "localhost:11434";
/// let prompt = "Hello ollama.";
/// let model = None; // Use default model "phi3"
///
/// match send_request(server, prompt, model) {
///     Ok(response) => println!("Response: {}", response),
///     Err(e) => eprintln!("Error: {}", e),
/// }
/// ```
pub fn send_request(server: &str, prompt: &str, model: Option<&str>) -> Result<String, Box<dyn Error>> {
    let model = model.unwrap_or("phi3");
    let mut stream = TcpStream::connect(server)?;

    let request_body = format!(
        r#"{{
            "model": "{}",
            "prompt": "{}",
            "options": {{
                "temperature": 0.2,
                "repeat_penalty": 1.5,
                "top_k": 25,
                "top_p": 0.25
            }}
        }}"#,
        model, prompt
    );

    let request = format!(
        "POST /api/generate HTTP/1.1\r\n\
         Host: localhost\r\n\
         Content-Type: application/json\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         \r\n\
         {}",
        request_body.len(),
        request_body
    );

    stream.write_all(request.as_bytes())?;
    stream.flush()?;

    let mut response = Vec::new();
    stream.read_to_end(&mut response)?;

    let response_text = String::from_utf8_lossy(&response);
    let mut full_response = String::new();
    let mut in_body = false;

    for line in response_text.lines() {
        if line.is_empty() {
            in_body = true;
            continue;
        }

        if in_body {
            full_response.push_str(line);
            full_response.push('\n');
        }
    }

    let mut final_response = String::new();
    for line in full_response.lines() {
        if let Some(start) = line.find(r#""response":"#) {
            let response_start = start + r#""response":"#.len() + 1;
            let response_part = &line[response_start..];
            if let Some(end) = response_part.find("\",\"") {
                final_response.push_str(&response_part[..end]);
            }
        }
    }

    if final_response.is_empty() {
        Ok("No response received.".to_string())
    } else {
        Ok(final_response.trim().to_string())
    }
}