Merge pull request #2 from douglasmakey/feature/use-rust-genai

feat: add rust-genai.

Cargo.toml

@@ -30,3 +30,4 @@ regex = "1.10.4"
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.117"
 tokio = { version = "1.38.0", features = ["full"] }
+genai = "=0.1.1"

README.md

@@ -2,7 +2,7 @@
 
 ![](./assets/shelldon.jpeg)
 
-Shelldon is a command-line tool written in Rust. It provides a set of utilities for executing shell commands, managing prompts, and interacting with OpenAI GPT.
+Shelldon is a command-line tool written in Rust. It provides utilities for executing shell commands, managing prompts, and interacting with multiple LLMs.
 
 Yes, another CLI with GPT features. Shelldon is not intended to be a full GPT client from the terminal; there are a couple of CLIs much better for that and also a lot of applications and even the OpenAI ChatGPT apps. Shelldon is to solve some personal use cases and it is very useful for me; I hope it could be useful for you too. Also, I made it to have fun playing with Rust!
 
@@ -32,18 +32,24 @@ cargo build --release
 
 ## Usage
 
-To use Shelldon, you’ll need to set your OpenAI token. You can do this by setting an environment variable. Here’s how you can set it in your terminal:
+Shelldon supports different AI providers such as Ollama, OpenAI, Gemini, Anthropic, and Cohere. You can control which provider to use with the `--model` flag. For example, `--model claude-3-haiku-20240307` or `--model gemini-1.5-flash-latest`. By default, Shelldon uses `gpt-4o` as the model.
+
+To use Shelldon, you need to set your API keys for the mentioned providers. You can do this by setting an environment variable. Here’s how to set it in your terminal:
 
 ```sh
-export OPENAI_API_KEY="your-openai-api-key"
+export OPENAI_API_KEY="api-key"
+export ANTHROPIC_API_KEY="api-key"
+export COHERE_API_KEY="api-key"
+export GEMINI_API_KEY="api-key"
+
 ```
 
 Shelldon allows you to integrate GPT features into your shell commands easily. Here are some examples to get you started:
 
 ### Running Shell Commands
 
 ```sh
-$ shelldon exec "Show all the graphics ports for the Vagrant machine using Libvirt."
+$ shelldon exec "Show all the graphics ports for the Vagrant machine using Libvirt." --model gpt-4o
 Command to execute: vagrant ssh -c "virsh list --all | grep vagrant | awk '{print \$1}' | xargs -I {} virsh domdisplay {}"
 ? [R]un, [M]odify, [C]opy, [A]bort › 
 ```
@@ -188,7 +194,7 @@ So the ability to handle dynamic prompts with args and use them makes Shelldon a
 - [ ] Improve error handling.
 - [ ] Add default prompts.
 - [ ] Implement OpenAI functions?
-- [ ] Implement Ollama? Maybe in the future. Do you need it?
+- [X] Implement Ollama? Maybe in the future. Do you need it?
 
 ## Contributing
 

src/backend/genai.rs

@@ -0,0 +1,59 @@
+use crate::processor::CompletionGenerator;
+use crate::{Error, Result};
+use async_stream::stream;
+use futures::{stream::LocalBoxStream, StreamExt};
+use genai::{
+    chat::{ChatMessage, ChatRequest, ChatStreamEvent, StreamChunk},
+    client::Client,
+};
+
+pub struct GenAI {
+    client: Client,
+}
+
+impl GenAI {
+    pub fn new() -> Self {
+        Self {
+            client: Client::default(),
+        }
+    }
+}
+
+impl CompletionGenerator for GenAI {
+    async fn generate_completion(
+        &self,
+        model: &str,
+        _temperature: f32,
+        prompt: &str,
+        input: &str,
+    ) -> crate::Result<String> {
+        let req = ChatRequest::new(vec![ChatMessage::system(prompt), ChatMessage::user(input)]);
+        let resp = self.client.exec_chat(model, req.clone(), None).await?;
+        resp.content.ok_or(Error::EmptyResponse)
+    }
+
+    async fn stream_completion(
+        &self,
+        model: &str,
+        _temperature: f32,
+        prompt: &str,
+        input: &str,
+    ) -> Result<LocalBoxStream<String>> {
+        let req = ChatRequest::new(vec![ChatMessage::system(prompt), ChatMessage::user(input)]);
+        let resp = self
+            .client
+            .exec_chat_stream(model, req.clone(), None)
+            .await?;
+
+        let async_stream = stream! {
+                let mut stream = resp.stream;
+                while let Some(Ok(stream_event)) = stream.next().await {
+                    if let ChatStreamEvent::Chunk(StreamChunk { content }) = stream_event {
+                        yield content;
+                    }
+                };
+        };
+
+        Ok(Box::pin(async_stream))
+    }
+}

src/backend/mod.rs

@@ -1 +1,2 @@
+pub mod genai;
 pub mod openai;

src/backend/openai.rs

@@ -9,7 +9,7 @@ use async_openai::{
     Client,
 };
 use async_stream::stream;
-use futures::stream::BoxStream;
+use futures::stream::LocalBoxStream;
 
 pub struct OpenAI {
     client: Client<OpenAIConfig>,
@@ -70,7 +70,7 @@ impl CompletionGenerator for OpenAI {
         temperature: f32,
         prompt: &str,
         input: &str,
-    ) -> Result<BoxStream<String>> {
+    ) -> Result<LocalBoxStream<String>> {
         let messages = [
             ChatCompletionRequestSystemMessageArgs::default()
                 .content(prompt)

src/command/ask.rs

@@ -1,6 +1,6 @@
 use super::CommonArgs;
 use crate::{
-    backend::openai::OpenAI,
+    backend::genai::GenAI,
     command::{parse_prompt, read_input},
     config::Config,
     processor::CompletionProcessor,
@@ -18,7 +18,7 @@ pub struct AskArgs {
 }
 
 pub async fn handle_ask(config: Config, args: AskArgs) -> Result<()> {
-    let processor = CompletionProcessor::new(OpenAI::new()?);
+    let processor = CompletionProcessor::new(GenAI::new());
     let input = read_input(&args.common.input)?;
     let prompt = parse_prompt(config, args.common.prompt, args.common.set, "")?;
     let mut completion = processor

src/command/exec.rs

@@ -1,6 +1,6 @@
 use super::{parse_prompt, read_input, CommonArgs};
 use crate::{
-    backend::openai::OpenAI,
+    backend::genai::GenAI,
     config::Config,
     processor::CompletionProcessor,
     system::{self, copy_to_clipboard, run_cmd},
@@ -36,7 +36,7 @@ pub struct ExecArgs {
 }
 
 pub async fn handle_exec(config: Config, args: ExecArgs) -> Result<()> {
-    let processor = CompletionProcessor::new(OpenAI::new()?);
+    let processor = CompletionProcessor::new(GenAI::new());
     let input = read_input(&args.common.input)?;
     let default_prompt = SHELL_PROMPT
         .replace("{shell}", &system::get_current_shell())

src/error.rs

@@ -12,6 +12,8 @@ pub enum Error {
     CommandFailed { command: String },
     #[display(fmt = "API key not set")]
     APIKeyNotSet,
+    #[display(fmt = "Empty response")]
+    EmptyResponse,
 
     #[from]
     OpenAI(async_openai::error::OpenAIError),
@@ -23,4 +25,6 @@ pub enum Error {
     Serde(serde_json::Error),
     #[from]
     Dialoguer(dialoguer::Error),
+    #[from]
+    GenAI(genai::Error),
 }

src/processor.rs

@@ -1,5 +1,5 @@
 use crate::Result;
-use futures::stream::BoxStream;
+use futures::stream::LocalBoxStream;
 
 pub trait CompletionGenerator {
     async fn generate_completion(
@@ -16,7 +16,7 @@ pub trait CompletionGenerator {
         temperature: f32,
         prompt: &str,
         input: &str,
-    ) -> Result<BoxStream<String>>;
+    ) -> Result<LocalBoxStream<String>>;
 }
 
 pub struct CompletionProcessor<T: CompletionGenerator> {
@@ -48,7 +48,7 @@ impl<T: CompletionGenerator> CompletionProcessor<T> {
         input: &str,
         model: &str,
         temperature: f32,
-    ) -> Result<BoxStream<String>> {
+    ) -> Result<LocalBoxStream<String>> {
         self.generator
             .stream_completion(model, temperature, prompt, input)
             .await