From 4b7f9d8e0d561c36d085981926a72f9f3b86372c Mon Sep 17 00:00:00 2001 From: wiseaidev Date: Mon, 4 May 2026 20:34:23 +0300 Subject: [PATCH] fix: drop deprecated models && use latest versions available. Signed-off-by: wiseaidev --- .circleci/config.yml | 2 +- LICENSE | 2 +- README.md | 75 +++++++++++++------------ examples/axum/README.md | 2 +- examples/basic.ipynb | 112 +++++++++++++++++++++++++++++++------- examples/rocket/README.md | 2 +- src/chat.rs | 14 +++++ src/cli.rs | 2 +- src/client.rs | 23 +++++--- src/embed.rs | 18 +++++- src/imagen.rs | 71 ++++++++++++++---------- src/lib.rs | 14 +++++ src/main.rs | 56 +++++++++++-------- src/messages.rs | 25 +++++++++ src/models.rs | 88 ++++++++++++++++-------------- src/requests.rs | 33 +++++++++++ src/responses.rs | 27 +++++++++ src/stream.rs | 14 +++++ src/tokens.rs | 14 +++++ src/traits.rs | 14 +++++ src/tts.rs | 14 +++++ src/tui.rs | 26 ++++++--- src/utils.rs | 14 +++++ src/vidgen.rs | 14 +++++ src/vision.rs | 16 +++++- tests/chat.rs | 46 +++++++++++----- tests/embed.rs | 32 ++++++++--- tests/imagen.rs | 32 ++++++++--- 28 files changed, 592 insertions(+), 210 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 5e62e44..a2fd8f4 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3,7 +3,7 @@ version: 2.1 aliases: - &rust_container docker: - - image: cimg/rust:1.86.0 + - image: cimg/rust:1.89.0 jobs: testing: <<: *rust_container diff --git a/LICENSE b/LICENSE index 14322be..4cd680f 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2025 Kevin RS Foundation +Copyright (c) 2026 Wise AI Foundation Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 7fff778..01a3fd9 100755 --- a/README.md +++ b/README.md @@ -6,12 +6,12 @@ [![docs](https://docs.rs/gems/badge.svg)](https://docs.rs/gems/) [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE) -| 🐧 Linux `(Recommended)` | 🪟 Windows | -| :------: | :--------: | -| ![gems-demo](https://github.com/user-attachments/assets/c446c29f-d4c8-4ee0-9e3d-951310e2b972) | ![gems-demo](https://github.com/user-attachments/assets/e942d1ad-7df6-4532-b22f-d4c586e64c8a) | -| [Download Executable File](https://github.com/kevin-rs/gems/releases/download/v0.1.4/gems) | [Download `.exe` File](https://github.com/kevin-rs/gems/releases/download/v0.1.4/gems.exe) | -| `cargo install gems --all-features` | `cargo install gems --all-features` | -| `gems -h` | `gems -h` | +| 🐧 Linux `(Recommended)` | 🪟 Windows | +| :--------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------: | +| ![gems-demo](https://github.com/user-attachments/assets/c446c29f-d4c8-4ee0-9e3d-951310e2b972) | ![gems-demo](https://github.com/user-attachments/assets/e942d1ad-7df6-4532-b22f-d4c586e64c8a) | +| [Download Executable File](https://github.com/wiseaidotdev/gems/releases/download/v0.1.4/gems) | [Download `.exe` File](https://github.com/wiseaidotdev/gems/releases/download/v0.1.4/gems.exe) | +| `cargo install gems --all-features` | `cargo install gems --all-features` | +| `gems -h` | `gems -h` | @@ -143,33 +143,30 @@ gems ## 🎨 Options -| Option | Description | -|--------------------------|----------------------------------------------------------| -| `` | TUI mode. | -| `--api-key` | Specify the API key for accessing the Gemini API. | -| `--model` | Specify the model to use for generating content. | - +| Option | Description | +| ----------- | ------------------------------------------------- | +| `` | TUI mode. | +| `--api-key` | Specify the API key for accessing the Gemini API. | +| `--model` | Specify the model to use for generating content. | ## 🛠 Subcommands -| Subcommand | Description | -|-------------------------|----------------------------------------------------------| -| `generate` | Generate creative content. | -| `vision` | Analyze an image and generate content from text. | -| `stream` | Stream the generation of content. | -| `imagen` | Generate an image. | -| `vidgen` | Generate a video. | -| `tts` | Text to speech. | -| `count` | Count the number of tokens in a text. | -| `embed` | Embed content into a specified model. | -| `batch` | Batch embed multiple contents. | -| `info` | Get information about the current model. | -| `list` | List available models. | - +| Subcommand | Description | +| ---------- | ------------------------------------------------ | +| `generate` | Generate creative content. | +| `vision` | Analyze an image and generate content from text. | +| `stream` | Stream the generation of content. | +| `imagen` | Generate an image. | +| `vidgen` | Generate a video. | +| `tts` | Text to speech. | +| `count` | Count the number of tokens in a text. | +| `embed` | Embed content into a specified model. | +| `batch` | Batch embed multiple contents. | +| `info` | Get information about the current model. | +| `list` | List available models. | ## ✨ Usage as Dependency - 1. Add the `gems` crate: ```toml @@ -195,7 +192,7 @@ gems gemini_client.set_api_key("your-api-key".to_string()); let parameters = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![Message::User { content: Content::Text("Hello".to_string()), name: None, @@ -222,7 +219,7 @@ This repository contains a list of notebooks examples on how to use the sdk and 1. Clone the repository to your local machine: ```sh - git clone https://github.com/kevin-rs/gems.git + git clone https://github.com/wiseaidotdev/gems.git ``` 1. Install the required dependencies and libraries. Make sure you have [`Rust`](https://rustup.rs/), [`Jupyter Notebook`](https://jupyter.org/install), and [`evcxr_jupyter`](https://github.com/evcxr/evcxr/blob/main/evcxr_jupyter/README.md) installed on your system. @@ -230,13 +227,13 @@ This repository contains a list of notebooks examples on how to use the sdk and ```sh # Install a Rust toolchain (e.g. nightly): curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain nightly - + # Install Jupyter Notebook pip install notebook - + # Install evcxr_jupyter cargo install evcxr_jupyter - evcxr_jupyter --install + evcxr_jupyter --install ``` 1. Navigate to the cloned repository and build the project: @@ -254,16 +251,18 @@ This repository contains a list of notebooks examples on how to use the sdk and 1. Access the notebooks in your web browser by clicking on the notebook file you want to explore. -| ID | Example | Open on GitHub | Launch on Binder | Launch on Colab | -|----|---------------|-----------|:-------------|-------------| -| 1 | **Basic** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/basic.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/kevin-rs/gems/main?filepath=examples/basic.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kevin-rs/gems/blob/main/examples/basic.ipynb) | -| 2 | **Rocket** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/rocket.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/kevin-rs/gems/main?filepath=examples/rocket.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kevin-rs/gems/blob/main/examples/rocket.ipynb) | -| 3 | **Axum** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/axum.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/kevin-rs/gems/main?filepath=examples/axum.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kevin-rs/gems/blob/main/examples/axum.ipynb) | +| ID | Example | Open on GitHub | Launch on Binder | Launch on Colab | +| --- | ---------- | ----------------------------------------------------------------------------------------------------------------------- | :--------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| 1 | **Basic** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/basic.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/wiseaidotdev/gems/main?filepath=examples/basic.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wiseaidotdev/gems/blob/main/examples/basic.ipynb) | +| 2 | **Rocket** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/rocket.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/wiseaidotdev/gems/main?filepath=examples/rocket.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wiseaidotdev/gems/blob/main/examples/rocket.ipynb) | +| 3 | **Axum** | [![Github](https://img.shields.io/badge/launch-Github-181717.svg?logo=github&logoColor=white)](./examples/axum.ipynb) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/wiseaidotdev/gems/main?filepath=examples/axum.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wiseaidotdev/gems/blob/main/examples/axum.ipynb) | ## 🤝 Contributing -Contributions and feedback are welcome! If you'd like to contribute, report an issue, or suggest an enhancement, please engage with the project on [GitHub](https://github.com/kevin-rs/gems). Your contributions help improve this crate for the community. +Contributions and feedback are welcome! If you'd like to contribute, report an issue, or suggest an enhancement, please engage with the project on [GitHub](https://github.com/wiseaidotdev/gems). Your contributions help improve this crate for the community. ## 📄 License This project is licensed under the [MIT License](LICENSE). + +© 2026 Wise AI Foundation diff --git a/examples/axum/README.md b/examples/axum/README.md index 8180e54..307aed3 100644 --- a/examples/axum/README.md +++ b/examples/axum/README.md @@ -102,4 +102,4 @@ This endpoint embeds content based on the input text. ## License -📜 This project is licensed under the [MIT](LICENSE) license - see the [LICENSE](LICENSE) file for details. \ No newline at end of file +📜 This project is licensed under the [MIT](LICENSE) license - see the [LICENSE](LICENSE) file for details. diff --git a/examples/basic.ipynb b/examples/basic.ipynb index a949ec9..ff88d1c 100644 --- a/examples/basic.ipynb +++ b/examples/basic.ipynb @@ -14,7 +14,11 @@ "cell_type": "code", "execution_count": null, "id": "37160ef5", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "# This script sets up and spins up a Jupyter Notebook environment with a Rust kernel using Nix and IPC Proxy. \n", @@ -44,7 +48,11 @@ "cell_type": "code", "execution_count": 3, "id": "2c9d82d9", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ ":dep gems = { version = \"0.1.0\" }" @@ -70,7 +78,11 @@ "cell_type": "code", "execution_count": 4, "id": "953b8dc3", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use gems::Client;\n", @@ -91,7 +103,11 @@ "cell_type": "code", "execution_count": 5, "id": "8aca06b6", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "fn config() -> (String, String) {\n", @@ -119,9 +135,13 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "id": "e7a01cc5", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "ename": "Error", @@ -144,7 +164,7 @@ "gemini_client.set_api_key(config().0);\n", "\n", "let parameters = ChatBuilder::default()\n", - " .model(Model::Flash20)\n", + " .model(Model::Flash3Preview)\n", " .messages(vec![Message::User {\n", " content: Content::Text(\"Hello\".to_string()),\n", " name: None,\n", @@ -181,7 +201,11 @@ "cell_type": "code", "execution_count": null, "id": "e9bab3a3", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use gems::stream::StreamBuilder\n", @@ -191,7 +215,7 @@ "let input_text = \"Write a story about a magic backpack.\";\n", " \n", "let parameters = StreamBuilder::default()\n", - " .model(Model::Flash20)\n", + " .model(Model::Flash3Preview)\n", " .input(Message::User {\n", " content: Content::Text(input_text.to_string()),\n", " name: None,\n", @@ -228,7 +252,11 @@ "cell_type": "code", "execution_count": null, "id": "bdbd3a16", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use gems::tokens::TokenBuilder\n", @@ -263,7 +291,11 @@ "cell_type": "code", "execution_count": null, "id": "d6c20478", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use gems::embed::EmbeddingBuilder;\n", @@ -299,7 +331,11 @@ "cell_type": "code", "execution_count": null, "id": "19376863", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use gems::models::Model;\n", @@ -325,7 +361,11 @@ "cell_type": "code", "execution_count": null, "id": "a8804b3c", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "let models = gemini_client.models().list().await?;\n", @@ -354,7 +394,11 @@ "cell_type": "code", "execution_count": 2, "id": "b71d5b4c", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [], "source": [ "use std::process::{Command, Output, Stdio};\n", @@ -379,7 +423,11 @@ "cell_type": "code", "execution_count": 3, "id": "6a53a882", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "name": "stdout", @@ -493,7 +541,11 @@ "cell_type": "code", "execution_count": 4, "id": "6748c283", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "name": "stdout", @@ -536,7 +588,11 @@ "cell_type": "code", "execution_count": 5, "id": "120d257b", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "name": "stdout", @@ -600,7 +656,11 @@ "cell_type": "code", "execution_count": 6, "id": "f4a870cc", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "name": "stdout", @@ -643,7 +703,10 @@ "execution_count": 7, "id": "02d9b2a9", "metadata": { - "scrolled": true + "scrolled": true, + "vscode": { + "languageId": "rust" + } }, "outputs": [ { @@ -686,7 +749,11 @@ "cell_type": "code", "execution_count": 8, "id": "d5860743", - "metadata": {}, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, "outputs": [ { "name": "stdout", @@ -742,7 +809,10 @@ "execution_count": 9, "id": "d875ac0a", "metadata": { - "scrolled": true + "scrolled": true, + "vscode": { + "languageId": "rust" + } }, "outputs": [ { diff --git a/examples/rocket/README.md b/examples/rocket/README.md index b39251c..093efd6 100644 --- a/examples/rocket/README.md +++ b/examples/rocket/README.md @@ -102,4 +102,4 @@ This endpoint embeds content based on the input text. ## License -📜 This project is licensed under the [MIT](LICENSE) license - see the [LICENSE](LICENSE) file for details. \ No newline at end of file +📜 This project is licensed under the [MIT](LICENSE) license - see the [LICENSE](LICENSE) file for details. diff --git a/src/chat.rs b/src/chat.rs index 1d7222c..23d5599 100644 --- a/src/chat.rs +++ b/src/chat.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; @@ -60,3 +67,10 @@ impl Chats { } } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/cli.rs b/src/cli.rs index 8f63df6..8d706c2 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -90,7 +90,7 @@ EXAMPLES: TUI mode: gems -For more information, visit: github.com/kevin-rs/gems +For more information, visit: github.com/wiseaidotdev/gems "# )] pub struct Cli { diff --git a/src/client.rs b/src/client.rs index ad01dfb..d561db9 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::chat::Chats; use crate::embed::Embeddings; use crate::imagen::Images; @@ -59,14 +66,9 @@ impl CTrait for Client { let full_url = if endpoint == "models" { GEMINI_BASE_URL.to_string() } else if endpoint.is_empty() { - format!("{}/{}", GEMINI_BASE_URL, self.get_model().to_string()) + format!("{}/{}", GEMINI_BASE_URL, self.get_model()) } else { - format!( - "{}/{}:{}", - GEMINI_BASE_URL, - self.get_model().to_string(), - endpoint - ) + format!("{}/{}:{}", GEMINI_BASE_URL, self.get_model(), endpoint) }; let parsed_url = Url::parse_with_params(&full_url, &[("key", api_key)]).unwrap(); @@ -162,3 +164,10 @@ impl CBuilder { }) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/embed.rs b/src/embed.rs index 705d259..e9d193c 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; @@ -33,7 +40,7 @@ pub struct Embeddings { impl Embeddings { pub async fn create(&self, params: Embedding) -> Result { let request_body = GeminiEmbedRequest { - model: format!("models/{}", params.model.to_string()), + model: format!("models/{}", params.model), content: Content { parts: vec![params.input.to_part()], }, @@ -53,7 +60,7 @@ impl Embeddings { .input .into_iter() .map(|message| GeminiEmbedRequest { - model: format!("models/{}", params.model.to_string()), + model: format!("models/{}", params.model), content: Content { parts: vec![message.to_part()], }, @@ -71,3 +78,10 @@ impl Embeddings { Ok(res.json().await?) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/imagen.rs b/src/imagen.rs index daf7640..213bed0 100644 --- a/src/imagen.rs +++ b/src/imagen.rs @@ -1,15 +1,20 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; -use crate::requests::GenerationConfig; -use crate::requests::{Content, GeminiRequest}; -use crate::responses::ImagenResponse; +use crate::requests::{ImagenParameters, ImagenPrompt, ImagenRequest}; +use crate::responses::ImagenPredictResponse; use crate::traits::CTrait; -use crate::utils::extract_image_or_text; use anyhow::{anyhow, Result}; +use base64::{engine::general_purpose, Engine as _}; use derive_builder::Builder; use reqwest::Method; - #[derive(Clone)] pub struct Images { pub client: Client, @@ -26,40 +31,46 @@ pub struct ImageGen { impl Images { pub async fn generate(&self, params: ImageGen) -> Result> { - let content = Content { - parts: vec![params.input.to_part()], - }; - - let system_instruction = params.system.as_ref().map(|messages| Content { - parts: messages.iter().map(|msg| msg.to_part()).collect(), - }); - - let request_body = GeminiRequest { - model: params.model.to_string(), - contents: vec![content], - system_instruction, - config: Some(GenerationConfig { - response_modalities: vec!["Text".into(), "Image".into()], - }), + let request_body = ImagenRequest { + instances: vec![ImagenPrompt { + prompt: params.input.get_text(), + }], + parameters: ImagenParameters { + sample_count: 1, + aspect_ratio: Some("1:1".to_string()), + }, }; let req = self .client - .request(Method::POST, "generateContent")? + .request(Method::POST, "predict")? .json(&request_body); let res = req.send().await?; - let json: ImagenResponse = res.json().await?; + if !res.status().is_success() { + let status = res.status(); + let error_text = res.text().await?; + return Err(anyhow!("API Error ({}): {}", status, error_text)); + } + + let json: ImagenPredictResponse = res.json().await?; - let parts = json - .candidates - .ok_or_else(|| anyhow!("Missing candidates"))? + let prediction = json + .predictions .first() - .ok_or_else(|| anyhow!("No candidate response"))? - .content - .parts - .clone(); + .ok_or_else(|| anyhow!("No predictions in response"))?; - extract_image_or_text(&parts) + let image_data = general_purpose::STANDARD + .decode(&prediction.bytes_base64_encoded) + .map_err(|e| anyhow!("Failed to decode base64 image: {}", e))?; + + Ok(image_data) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/lib.rs b/src/lib.rs index d3015ee..58f8b7b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + #![cfg_attr(docsrs, feature(doc_auto_cfg))] #![doc = include_str!("../README.md")] @@ -23,3 +30,10 @@ pub mod cli; pub mod tui; pub use client::Client; + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/main.rs b/src/main.rs index a6b4821..d1b0d31 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use anyhow::Result; /// The main entry point of `gems`. @@ -35,26 +42,20 @@ async fn main() -> Result<()> { let args: Cli = Cli::parse(); - let api_key = if args.api_key.is_none() { - env::var("GEMINI_API_KEY").unwrap_or_default().to_owned() - } else { - args.api_key.unwrap().to_owned() - }; + let api_key = args + .api_key + .unwrap_or_else(|| env::var("GEMINI_API_KEY").unwrap_or_default()); - let model = if args.model.is_none() { - env::var("GEMINI_MODEL") - .unwrap_or("gemini-2.0-flash".to_string()) - .to_owned() - } else { - args.model.unwrap().to_owned() - }; + let model = args.model.unwrap_or_else(|| { + env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-3-flash-preview".to_string()) + }); let mut gemini_client = Client::builder().model(&model).build()?; gemini_client.set_api_key(api_key); match args.cmd { Some(Command::Generate(cmd)) => { let parameters = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![Message::User { content: Content::Text(cmd.text), name: None, @@ -66,7 +67,7 @@ async fn main() -> Result<()> { } Some(Command::Stream(cmd)) => { let parameters = StreamBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .input(Message::User { content: Content::Text(cmd.text), name: None, @@ -114,13 +115,13 @@ async fn main() -> Result<()> { } Some(Command::Embed(cmd)) => { let params = EmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(Message::User { content: Content::Text(cmd.text), name: None, }) .build()?; - gemini_client.set_model(Model::Embedding); + gemini_client.set_model(Model::Embedding001); let response = gemini_client.embeddings().create(params).await?; println!("Embed Content: {:?}", response); } @@ -134,11 +135,11 @@ async fn main() -> Result<()> { }) .collect(); let params = BatchEmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(texts) .build()?; - gemini_client.set_model(Model::Embedding); + gemini_client.set_model(Model::Embedding001); let response = gemini_client.embeddings().batch(params).await?; println!("Batch Embed Contents: {:?}", response); } @@ -173,14 +174,14 @@ async fn main() -> Result<()> { models.print(); } Some(Command::Imagen(cmd)) => { - gemini_client.set_model(Model::FlashExpImage); + gemini_client.set_model(Model::Imagen4); let params = ImageGenBuilder::default() .input(Message::User { content: Content::Text(cmd.text), name: None, }) - .model(Model::FlashExpImage) + .model(Model::Imagen4) .build() .unwrap(); @@ -189,10 +190,10 @@ async fn main() -> Result<()> { tokio::fs::write("output.png", &image_data).await?; } Some(Command::Vidgen(cmd)) => { - gemini_client.set_model(Model::Veo2); + gemini_client.set_model(Model::Veo31Preview); let params = VideoGenBuilder::default() - .model(Model::Veo2) + .model(Model::Veo31Preview) .input(Message::User { content: Content::Text(cmd.text), name: None, @@ -205,10 +206,10 @@ async fn main() -> Result<()> { tokio::fs::write("output.mp4", &bytes).await?; } Some(Command::Tts(cmd)) => { - gemini_client.set_model(Model::Tts); + gemini_client.set_model(Model::Tts31Preview); let params = TtsGenBuilder::default() - .model(Model::Tts) + .model(Model::Tts31Preview) .input(Message::User { content: Content::Text(cmd.text), name: None, @@ -228,3 +229,10 @@ async fn main() -> Result<()> { } Ok(()) } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/messages.rs b/src/messages.rs index 5d442a8..a923c81 100644 --- a/src/messages.rs +++ b/src/messages.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::requests::Part; #[derive(Debug, Clone, PartialEq)] @@ -50,4 +57,22 @@ impl Message { Message::Tool { content } => Part::text(content), } } + + pub fn get_text(&self) -> String { + match self { + Message::User { content, .. } + | Message::System { content, .. } + | Message::Developer { content, .. } => match content { + Content::Text(text) => text.clone(), + }, + Message::Tool { content } => content.clone(), + } + } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/models.rs b/src/models.rs index cfe723f..e6cdfc0 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::responses::ModelInfo; use crate::responses::ModelsResponse; use crate::traits::CTrait; @@ -9,39 +16,32 @@ use std::str::FromStr; #[derive(Debug, Clone, Default, PartialEq)] pub enum Model { - Pro25Preview, + Pro31Preview, #[default] - Flash20, - Flash20Lite, - Flash15, - Flash15_8B, - Pro15, - Embedding, - Imagen3, - Veo2, - Tts, - Flash20Live, - FlashExpImage, + Flash3Preview, + Flash31LitePreview, + Embedding001, + Imagen4, + Veo31Preview, + Tts31Preview, + Flash31LivePreview, + Flash31ImagePreview, } -#[allow(clippy::to_string_trait_impl)] -impl ToString for Model { - fn to_string(&self) -> String { - match self { - Model::Pro25Preview => "gemini-2.5-pro-preview-03-25", - Model::Flash20 => "gemini-2.0-flash", - Model::Flash20Lite => "gemini-2.0-flash-lite", - Model::Flash15 => "gemini-1.5-flash", - Model::Flash15_8B => "gemini-1.5-flash-8b", - Model::Pro15 => "gemini-1.5-pro", - Model::Embedding => "text-embedding-004", - Model::Imagen3 => "imagen-3.0-generate-002", - Model::Veo2 => "veo-2.0-generate-001", - Model::Tts => "gemini-2.5-flash-preview-tts", - Model::Flash20Live => "gemini-2.0-flash-live-001", - Model::FlashExpImage => "gemini-2.0-flash-exp-image-generation", - } - .to_string() +impl std::fmt::Display for Model { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + Model::Pro31Preview => "gemini-3.1-pro-preview", + Model::Flash3Preview => "gemini-3-flash-preview", + Model::Flash31LitePreview => "gemini-3.1-flash-lite-preview", + Model::Embedding001 => "gemini-embedding-001", + Model::Imagen4 => "imagen-4.0-generate-001", + Model::Veo31Preview => "veo-3.1-generate-preview", + Model::Tts31Preview => "gemini-3.1-flash-tts-preview", + Model::Flash31LivePreview => "gemini-3.1-flash-live-preview", + Model::Flash31ImagePreview => "gemini-3.1-flash-image-preview", + }; + write!(f, "{}", s) } } @@ -50,18 +50,15 @@ impl FromStr for Model { fn from_str(s: &str) -> Result { match s { - "gemini-2.5-pro-preview-03-25" => Ok(Model::Pro25Preview), - "gemini-2.0-flash" => Ok(Model::Flash20), - "gemini-2.0-flash-lite" => Ok(Model::Flash20Lite), - "gemini-1.5-flash" => Ok(Model::Flash15), - "gemini-1.5-flash-8b" => Ok(Model::Flash15_8B), - "gemini-1.5-pro" => Ok(Model::Pro15), - "text-embedding-004" => Ok(Model::Embedding), - "imagen-3.0-generate-002" => Ok(Model::Imagen3), - "veo-2.0-generate-001" => Ok(Model::Veo2), - "gemini-2.5-flash-preview-tts" => Ok(Model::Tts), - "gemini-2.0-flash-live-001" => Ok(Model::Flash20Live), - "gemini-2.0-flash-exp-image-generation" => Ok(Model::Flash20Live), + "gemini-3.1-pro-preview" => Ok(Model::Pro31Preview), + "gemini-3-flash-preview" => Ok(Model::Flash3Preview), + "gemini-3.1-flash-lite-preview" => Ok(Model::Flash31LitePreview), + "gemini-embedding-001" => Ok(Model::Embedding001), + "imagen-4.0-generate-001" => Ok(Model::Imagen4), + "veo-3.1-generate-preview" => Ok(Model::Veo31Preview), + "gemini-3.1-flash-tts-preview" => Ok(Model::Tts31Preview), + "gemini-3.1-flash-live-preview" => Ok(Model::Flash31LivePreview), + "gemini-3.1-flash-image-preview" => Ok(Model::Flash31ImagePreview), _ => Err(anyhow!("Unknown model: {}", s)), } } @@ -92,3 +89,10 @@ impl Models { Ok(res.json().await?) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/requests.rs b/src/requests.rs index 2acc49b..b06c6d2 100644 --- a/src/requests.rs +++ b/src/requests.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use serde::{Deserialize, Serialize}; /// Request structure for content generation. @@ -144,3 +151,29 @@ pub struct PrebuiltVoiceConfig { #[serde(rename = "voiceName")] pub voice_name: String, } + +#[derive(Debug, Serialize)] +pub struct ImagenRequest { + pub instances: Vec, + pub parameters: ImagenParameters, +} + +#[derive(Debug, Serialize)] +pub struct ImagenPrompt { + pub prompt: String, +} + +#[derive(Debug, Serialize)] +pub struct ImagenParameters { + #[serde(rename = "sampleCount")] + pub sample_count: u32, + #[serde(rename = "aspectRatio", skip_serializing_if = "Option::is_none")] + pub aspect_ratio: Option, +} + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/responses.rs b/src/responses.rs index f4f5518..68387f1 100644 --- a/src/responses.rs +++ b/src/responses.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::requests::Candidate as ReqCandidate; use serde::{Deserialize, Serialize}; @@ -252,3 +259,23 @@ pub struct InlineData { pub mime_type: String, pub data: String, } + +#[derive(Debug, Deserialize)] +pub struct ImagenPredictResponse { + pub predictions: Vec, +} + +#[derive(Debug, Deserialize)] +pub struct ImagenPrediction { + #[serde(rename = "mimeType")] + pub mime_type: String, + #[serde(rename = "bytesBase64Encoded")] + pub bytes_base64_encoded: String, +} + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/stream.rs b/src/stream.rs index 28bdfe0..47c9bea 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; @@ -47,3 +54,10 @@ impl Streaming { Ok(res) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/tokens.rs b/src/tokens.rs index d5fe906..80f046d 100644 --- a/src/tokens.rs +++ b/src/tokens.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::messages::Message; use crate::models::Model; use crate::requests::Content; @@ -48,3 +55,10 @@ impl Tokens { Ok(json["totalTokens"].as_u64().unwrap_or(0) as usize) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/traits.rs b/src/traits.rs index a696b3b..bd759ba 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::chat::Chats; use crate::embed::Embeddings; use crate::imagen::Images; @@ -28,3 +35,10 @@ pub trait CTrait { fn videos(&self) -> Videos; fn tts(&self) -> Tts; } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/tts.rs b/src/tts.rs index 0562a08..5370728 100644 --- a/src/tts.rs +++ b/src/tts.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; @@ -103,3 +110,10 @@ impl Tts { .map_err(|e| anyhow!("Failed to decode audio: {}", e)) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/tui.rs b/src/tui.rs index edbd974..ea88c10 100644 --- a/src/tui.rs +++ b/src/tui.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use anyhow::Result; use crossterm::{ event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode}, @@ -95,7 +102,7 @@ impl Default for App { input_mode: InputMode::Normal, selected_tab: Tab::Settings, api_key: Input::default(), - selected_model: Input::new("gemini-2.0-flash".to_string()), + selected_model: Input::new("gemini-3-flash-preview".to_string()), user_input: Input::default(), chat_history: vec![], current_input: None, @@ -152,10 +159,8 @@ async fn run_app(terminal: &mut Terminal) -> Result<()> { KeyCode::Left | KeyCode::Char('a') => { app.selected_tab = app.selected_tab.previous() } - KeyCode::Up => { - if app.scroll_chat > 0 { - app.scroll_chat -= 1; - } + KeyCode::Up if app.scroll_chat > 0 => { + app.scroll_chat -= 1; } KeyCode::Down => { app.scroll_chat += 1; @@ -181,7 +186,7 @@ async fn run_app(terminal: &mut Terminal) -> Result<()> { app.chat_history.push(user_msg); let parameters = StreamBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .input(Message::User { content: Content::Text(msg), name: None, @@ -463,7 +468,7 @@ fn ui(f: &mut Frame, app: &mut App) { ) .centered(); - let bottom_footer = Line::raw("© Kevin RS Foundation") + let bottom_footer = Line::raw("© Wise AI Foundation") .fg(Color::LightGreen) .bg(tailwind::SLATE.c700) .bold() @@ -472,3 +477,10 @@ fn ui(f: &mut Frame, app: &mut App) { f.render_widget(top_footer, footer_layout[0]); f.render_widget(bottom_footer, footer_layout[2]); } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/utils.rs b/src/utils.rs index 676f591..3d116c4 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::responses::Part; use anyhow::{anyhow, Result}; use base64::{engine::general_purpose::STANDARD, Engine as _}; @@ -115,3 +122,10 @@ pub fn extract_image_or_text(parts: &[Part]) -> Result> { Err(anyhow!("No image or text found in response")) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/vidgen.rs b/src/vidgen.rs index 664def6..88b78ab 100644 --- a/src/vidgen.rs +++ b/src/vidgen.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::client::Client; use crate::messages::Message; use crate::models::Model; @@ -103,3 +110,10 @@ impl Videos { Err(anyhow!("Timed out waiting for video generation")) } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/src/vision.rs b/src/vision.rs index aa11bd2..100b64f 100644 --- a/src/vision.rs +++ b/src/vision.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use crate::messages::Content; use crate::messages::Message; use crate::models::Model; @@ -50,7 +57,7 @@ impl Visions { })); let request_body = GeminiRequest { - model: Model::Pro25Preview.to_string(), + model: Model::Pro31Preview.to_string(), contents: vec![crate::requests::Content { parts: vec![input_part, image_part], }], @@ -76,3 +83,10 @@ impl Visions { } } } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/tests/chat.rs b/tests/chat.rs index 125e3c6..0f706b9 100644 --- a/tests/chat.rs +++ b/tests/chat.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use anyhow::Result; use gems::chat::ChatBuilder; use gems::messages::Content; @@ -9,7 +16,7 @@ use gems::Client; #[test] fn test_build_with_required() { let chat = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![Message::User { content: Content::Text("Hello".into()), name: None, @@ -17,7 +24,7 @@ fn test_build_with_required() { .build() .unwrap(); - assert_eq!(chat.model, Model::Flash20); + assert_eq!(chat.model, Model::Flash3Preview); assert_eq!(chat.messages.len(), 1); assert!(chat.system.is_none()); } @@ -35,13 +42,13 @@ fn test_build_with_system() { }; let chat = ChatBuilder::default() - .model(Model::Pro15) + .model(Model::Pro31Preview) .messages(vec![user_message.clone()]) .system(vec![system_message.clone()]) .build() .unwrap(); - assert_eq!(chat.model, Model::Pro15); + assert_eq!(chat.model, Model::Pro31Preview); assert_eq!(chat.messages.len(), 1); assert_eq!(chat.system.as_ref().unwrap().len(), 1); } @@ -49,7 +56,7 @@ fn test_build_with_system() { #[test] fn test_build_empty_messages() { let result = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![]) .build(); @@ -85,7 +92,7 @@ fn test_build_clone() { }; let chat = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![user_message.clone()]) .build() .unwrap(); @@ -97,7 +104,7 @@ fn test_build_clone() { #[tokio::test] async fn test_required() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Flash20.to_string()) + .model(&Model::Flash3Preview.to_string()) .build()?; gemini_client.set_api_key( @@ -107,7 +114,7 @@ async fn test_required() -> Result<()> { ); let chat = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![Message::User { content: Content::Text("Hello".into()), name: None, @@ -128,7 +135,9 @@ async fn test_required() -> Result<()> { #[tokio::test] #[ignore] async fn test_with_system() -> Result<()> { - let gemini_client = Client::builder().model(&Model::Pro15.to_string()).build()?; + let gemini_client = Client::builder() + .model(&Model::Pro31Preview.to_string()) + .build()?; gemini_client.set_api_key( std::env::var("GEMINI_API_KEY") @@ -147,7 +156,7 @@ async fn test_with_system() -> Result<()> { }; let chat = ChatBuilder::default() - .model(Model::Pro15) + .model(Model::Pro31Preview) .messages(vec![user_message.clone()]) .system(vec![system_message.clone()]) .build() @@ -171,7 +180,7 @@ async fn test_with_system() -> Result<()> { #[tokio::test] async fn test_empty_msgs() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Flash20.to_string()) + .model(&Model::Flash3Preview.to_string()) .build()?; gemini_client.set_api_key( @@ -181,7 +190,7 @@ async fn test_empty_msgs() -> Result<()> { ); let chat = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![]) .build() .unwrap(); @@ -199,7 +208,7 @@ async fn test_empty_msgs() -> Result<()> { #[tokio::test] async fn test_default() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Flash20.to_string()) + .model(&Model::Flash3Preview.to_string()) .build()?; gemini_client.set_api_key( @@ -226,7 +235,7 @@ async fn test_default() -> Result<()> { #[tokio::test] async fn test_clone() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Flash20.to_string()) + .model(&Model::Flash3Preview.to_string()) .build()?; gemini_client.set_api_key( @@ -241,7 +250,7 @@ async fn test_clone() -> Result<()> { }; let chat = ChatBuilder::default() - .model(Model::Flash20) + .model(Model::Flash3Preview) .messages(vec![user_message.clone()]) .build() .unwrap(); @@ -260,3 +269,10 @@ async fn test_clone() -> Result<()> { Ok(()) } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/tests/embed.rs b/tests/embed.rs index ab7a25c..b0a98c4 100644 --- a/tests/embed.rs +++ b/tests/embed.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use anyhow::Result; use gems::embed::{BatchEmbeddingBuilder, EmbeddingBuilder}; use gems::messages::{Content, Message}; @@ -13,12 +20,12 @@ fn test_embed_build_required() { }; let embed = EmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(message.clone()) .build() .unwrap(); - assert_eq!(embed.model, Model::Embedding); + assert_eq!(embed.model, Model::Embedding001); assert_eq!(embed.input, message); } @@ -36,19 +43,19 @@ fn test_batch_embed_build() { ]; let batch = BatchEmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(inputs.clone()) .build() .unwrap(); - assert_eq!(batch.model, Model::Embedding); + assert_eq!(batch.model, Model::Embedding001); assert_eq!(batch.input, inputs); } #[tokio::test] async fn test_embed_create() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Embedding.to_string()) + .model(&Model::Embedding001.to_string()) .build()?; gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default()); @@ -59,7 +66,7 @@ async fn test_embed_create() -> Result<()> { }; let params = EmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(input) .build()?; @@ -84,7 +91,7 @@ async fn test_embed_create() -> Result<()> { ); } - assert_eq!(embeddings.len(), 768, "Unexpected embedding vector size"); + assert_eq!(embeddings.len(), 3072, "Unexpected embedding vector size"); Ok(()) } @@ -92,7 +99,7 @@ async fn test_embed_create() -> Result<()> { #[tokio::test] async fn test_embed_batch() -> Result<()> { let gemini_client = Client::builder() - .model(&Model::Embedding.to_string()) + .model(&Model::Embedding001.to_string()) .build()?; gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default()); @@ -109,7 +116,7 @@ async fn test_embed_batch() -> Result<()> { ]; let params = BatchEmbeddingBuilder::default() - .model(Model::Embedding) + .model(Model::Embedding001) .input(inputs) .build()?; @@ -123,3 +130,10 @@ async fn test_embed_batch() -> Result<()> { Ok(()) } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. diff --git a/tests/imagen.rs b/tests/imagen.rs index 5f9a427..d26deeb 100644 --- a/tests/imagen.rs +++ b/tests/imagen.rs @@ -1,3 +1,10 @@ +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + use anyhow::Result; use gems::imagen::ImageGenBuilder; use gems::messages::{Content, Message}; @@ -13,12 +20,12 @@ fn test_imagegen_build_required() { }; let imagegen = ImageGenBuilder::default() - .model(Model::FlashExpImage) + .model(Model::Imagen4) .input(user_msg.clone()) .build() .unwrap(); - assert_eq!(imagegen.model, Model::FlashExpImage); + assert_eq!(imagegen.model, Model::Imagen4); assert_eq!(imagegen.input, user_msg); assert!(imagegen.system.is_none()); } @@ -36,22 +43,23 @@ fn test_imagegen_build_with_system() { }; let imagegen = ImageGenBuilder::default() - .model(Model::FlashExpImage) + .model(Model::Imagen4) .input(user_msg.clone()) .system(vec![system_msg.clone()]) .build() .unwrap(); - assert_eq!(imagegen.model, Model::FlashExpImage); + assert_eq!(imagegen.model, Model::Imagen4); assert_eq!(imagegen.input, user_msg); assert_eq!(imagegen.system.unwrap().len(), 1); } #[tokio::test] +#[ignore] async fn test_imagegen_generate_basic() -> Result<()> { let mut gemini_client = Client::builder().build()?; - gemini_client.set_model(Model::FlashExpImage); + gemini_client.set_model(Model::Imagen4); gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default()); let input = Message::User { @@ -60,7 +68,7 @@ async fn test_imagegen_generate_basic() -> Result<()> { }; let params = ImageGenBuilder::default() - .model(Model::FlashExpImage) + .model(Model::Imagen4) .input(input) .build()?; @@ -71,13 +79,12 @@ async fn test_imagegen_generate_basic() -> Result<()> { Ok(()) } -// TODO: Research why Gemini doesn't allow system prompts in image gen #[tokio::test] #[ignore] async fn test_imagegen_generate_with_system() -> Result<()> { let mut gemini_client = Client::builder().build()?; - gemini_client.set_model(Model::FlashExpImage); + gemini_client.set_model(Model::Imagen4); gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default()); let input = Message::User { @@ -91,7 +98,7 @@ async fn test_imagegen_generate_with_system() -> Result<()> { }; let params = ImageGenBuilder::default() - .model(Model::FlashExpImage) + .model(Model::Imagen4) .input(input) .system(vec![system]) .build()?; @@ -105,3 +112,10 @@ async fn test_imagegen_generate_with_system() -> Result<()> { Ok(()) } + +// Copyright 2026 Mahmoud Harmouch. +// +// Licensed under the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms.