ruchat 0.1.2

ollama/chroma command-line AI chat tool
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

use crate::error::RuChatError;
use crate::io::Io;
use crate::ollama::model::get_name;
use crate::options::get_options;
use clap::Parser;
use ollama_rs::{Ollama, generation::completion::request::GenerationRequest};
use tokio_stream::StreamExt;

/// Command-line arguments for piping a question to a model.
///
/// This struct defines the arguments required to pipe a question
/// to a model, including model details and configuration options.
#[derive(Parser, Debug, Clone, Default, PartialEq)]
pub struct PipeArgs {
    /// Initial model to (down)load and use.
    #[clap(short, long, default_value = "qwen2.5-coder:14b")]
    pub(crate) model: Option<String>,

    /// Path to a JSON file to amend default generation options.
    #[clap(short, long)]
    pub(crate) config: Option<String>,

    /// Specify the model using a positional argument.
    pub(crate) positional_model: Option<String>,
}

/// The pipe command handles prompted questions with context using a model.
///
/// This function connects to a model using the provided arguments,
/// generates a response to the specified prompt, and outputs the response.
///
/// # Parameters
///
/// - `ollama`: The Ollama client for generating responses.
/// - `args`: The command-line arguments for the pipe operation.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub(crate) async fn pipe(ollama: Ollama, args: &PipeArgs) -> Result<(), RuChatError> {
    let mut cio = Io::new();
    let mut done = false;
    let mut options = get_options(&args.config).await?;
    let mut model_name = match args.model.as_deref().or(args.positional_model.as_deref()) {
        Some(model) if !model.is_empty() => get_name(&ollama, model).await?,
        _ => {
            return Err(RuChatError::ModelError(
                "Model name is required".to_string(),
            ));
        }
    };
    while !done {
        let mut prompt = String::new();
        while let Ok(line) = cio.read_line(false).await {
            match line.as_str() {
                "" => {
                    done = true;
                    break;
                }
                "---" | "***" | "___" => break,
                _ => prompt.push_str(&line),
            }
        }

        prompt.push_str("\nRespond in CommonMark format.\n");
        let request = GenerationRequest::new(model_name.clone(), prompt).options(options.clone());
        let mut stream = ollama.generate_stream(request).await?;
        while let Some(res) = stream.next().await {
            let responses = res?;
            for resp in responses {
                cio.write_line(&resp.response).await?;
            }
        }
    }
    Ok(())
}