feat(xai): initial xai (grok) implementation

README.md

@@ -39,8 +39,8 @@ Help us improve Rig by contributing to our [Feedback form](https://bit.ly/Rig-Fe
 - [What is Rig?](#what-is-rig)
 - [Table of contents](#table-of-contents)
 - [High-level features](#high-level-features)
-- [Installation](#installation)
-- [Simple example:](#simple-example)
+- [Get Started](#get-started)
+  - [Simple example:](#simple-example)
 - [Integrations](#integrations)
 
 ## High-level features
@@ -85,7 +85,7 @@ You can find more examples each crate's `examples` (ie. [`src/examples`](./src/e
 
 | Model Providers | Vector Stores |
 |:--------------:|:-------------:|
-| <br><img src="https://upload.wikimedia.org/wikipedia/commons/thumb/0/04/ChatGPT_logo.svg/1024px-ChatGPT_logo.svg.png" alt="ChatGPT logo" width="50em"> <picture><source media="(prefers-color-scheme: dark)" srcset="https://www.fahimai.com/wp-content/uploads/2024/06/Untitled-design-7.png"><source media="(prefers-color-scheme: light)" srcset="https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/Claude_Ai.svg/1024px-Claude_Ai.svg.png"><img src="https://www.fahimai.com/wp-content/uploads/2024/06/Untitled-design-7.png" alt="Claude Anthropic logo" width="50em"></picture> <br> <img src="https://cdn.sanity.io/images/rjtqmwfu/production/0adbf394439f4cd0ab8b5b3b6fe1da10c8099024-201x200.svg" alt="Cohere logo" width="50em"> <img src="https://logospng.org/download/google-gemini/google-gemini-1024.png" style="background-color: white; border-radius: 10px; padding: 5px 5px ; width: 3em;" alt="Gemini logo"> |  <br><img src="https://cdn.prod.website-files.com/6640cd28f51f13175e577c05/664e00a400e23f104ed2b6cd_3b3dd6e8-8a73-5879-84a9-a42d5b910c74.svg" alt="Mongo DB logo" width="50em"> <img src="https://upload.wikimedia.org/wikipedia/commons/e/e5/Neo4j-logo_color.png" alt="Neo4j logo" style="background-color: white; border-radius: 1em; padding: 1em 1em ; width: 4em;"><br><br><img src="https://cdn-images-1.medium.com/max/844/1*[email protected]" width="100em" alt="Lance DB logo"> |
+| <br><img src="https://upload.wikimedia.org/wikipedia/commons/thumb/0/04/ChatGPT_logo.svg/1024px-ChatGPT_logo.svg.png" alt="ChatGPT logo" width="50em"> <picture><source media="(prefers-color-scheme: dark)" srcset="https://www.fahimai.com/wp-content/uploads/2024/06/Untitled-design-7.png"><source media="(prefers-color-scheme: light)" srcset="https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/Claude_Ai.svg/1024px-Claude_Ai.svg.png"><img src="https://www.fahimai.com/wp-content/uploads/2024/06/Untitled-design-7.png" alt="Claude Anthropic logo" width="50em"></picture> <br> <img src="https://cdn.sanity.io/images/rjtqmwfu/production/0adbf394439f4cd0ab8b5b3b6fe1da10c8099024-201x200.svg" alt="Cohere logo" width="50em"> <img src="https://logospng.org/download/google-gemini/google-gemini-1024.png" style="background-color: white; border-radius: 10px; padding: 5px 5px ; width: 3em;" alt="Gemini logo"> <br> <img src="https://upload.wikimedia.org/wikipedia/commons/thumb/5/57/XAI-Logo.svg/512px-XAI-Logo.svg.png?20240912222841" style="background-color: white; border-radius: 10px; padding: 5px 5px ; width: 3em;" alt="xAI logo"> |  <br><img src="https://cdn.prod.website-files.com/6640cd28f51f13175e577c05/664e00a400e23f104ed2b6cd_3b3dd6e8-8a73-5879-84a9-a42d5b910c74.svg" alt="Mongo DB logo" width="50em"> <img src="https://upload.wikimedia.org/wikipedia/commons/e/e5/Neo4j-logo_color.png" alt="Neo4j logo" style="background-color: white; border-radius: 1em; padding: 1em 1em ; width: 4em;"><br><br><img src="https://cdn-images-1.medium.com/max/844/1*[email protected]" width="100em" alt="Lance DB logo"> |
 
 
 Vector stores are available as separate companion-crates:

rig-core/README.md

@@ -4,10 +4,12 @@ Rig is a Rust library for building LLM-powered applications that focuses on ergo
 More information about this crate can be found in the [crate documentation](https://docs.rs/rig-core/latest/rig/).
 ## Table of contents
 
-- [High-level features](#high-level-features)
-- [Installation](#)
-- [Simple Example](#simple-example)
-- [Integrations](#integrations)
+- [Rig](#rig)
+  - [Table of contents](#table-of-contents)
+  - [High-level features](#high-level-features)
+  - [Installation](#installation)
+  - [Simple example:](#simple-example)
+  - [Integrations](#integrations)
 
 ## High-level features
 - Full support for LLM completion and embedding workflows
@@ -48,5 +50,7 @@ Rig supports the following LLM providers natively:
 - OpenAI
 - Cohere
 - Google Gemini
+- xAI
+
 Additionally, Rig currently has the following integration sub-libraries:
 - MongoDB vector store: `rig-mongodb`

rig-core/examples/agent_with_grok.rs

@@ -0,0 +1,209 @@
+use std::env;
+
+use rig::{
+    agent::AgentBuilder,
+    completion::{Prompt, ToolDefinition},
+    loaders::FileLoader,
+    providers,
+    tool::Tool,
+};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+
+/// Runs 4 agents based on grok (dervived from the other examples)
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    println!("Running basic agent with grok");
+    basic().await?;
+
+    println!("\nRunning grok agent with tools");
+    tools().await?;
+
+    println!("\nRunning grok agent with loaders");
+    loaders().await?;
+
+    println!("\nRunning grok agent with context");
+    context().await?;
+
+    println!("\n\nAll agents ran successfully");
+    Ok(())
+}
+
+fn client() -> providers::xai::Client {
+    providers::xai::Client::new(&env::var("XAI_API_KEY").expect("XAI_API_KEY not set"))
+}
+
+/// Create a partial xAI agent (grok)
+fn partial_agent() -> AgentBuilder<providers::xai::completion::CompletionModel> {
+    let client = client();
+    client.agent(providers::xai::GROK_BETA)
+}
+
+/// Create an xAI agent (grok) with a preamble
+/// Based upon the `agent` example
+///
+/// This example creates a comedian agent with a preamble
+async fn basic() -> Result<(), anyhow::Error> {
+    let comedian_agent = partial_agent()
+        .preamble("You are a comedian here to entertain the user using humour and jokes.")
+        .build();
+
+    // Prompt the agent and print the response
+    let response = comedian_agent.prompt("Entertain me!").await?;
+    println!("{}", response);
+
+    Ok(())
+}
+
+/// Create an xAI agent (grok) with tools
+/// Based upon the `tools` example
+///
+/// This example creates a calculator agent with two tools: add and subtract
+async fn tools() -> Result<(), anyhow::Error> {
+    // Create agent with a single context prompt and two tools
+    let calculator_agent = partial_agent()
+        .preamble("You are a calculator here to help the user perform arithmetic operations. Use the tools provided to answer the user's question.")
+        .max_tokens(1024)
+        .tool(Adder)
+        .tool(Subtract)
+        .build();
+
+    // Prompt the agent and print the response
+    println!("Calculate 2 - 5");
+    println!(
+        "Calculator Agent: {}",
+        calculator_agent.prompt("Calculate 2 - 5").await?
+    );
+
+    Ok(())
+}
+
+/// Create an xAI agent (grok) with loaders
+/// Based upon the `loaders` example
+///
+/// This example loads in all the rust examples from the rig-core crate and uses them as\\
+///  context for the agent
+async fn loaders() -> Result<(), anyhow::Error> {
+    let model = client().completion_model(providers::xai::GROK_BETA);
+
+    // Load in all the rust examples
+    let examples = FileLoader::with_glob("rig-core/examples/*.rs")?
+        .read_with_path()
+        .ignore_errors()
+        .into_iter();
+
+    // Create an agent with multiple context documents
+    let agent = examples
+        .fold(AgentBuilder::new(model), |builder, (path, content)| {
+            builder.context(format!("Rust Example {:?}:\n{}", path, content).as_str())
+        })
+        .build();
+
+    // Prompt the agent and print the response
+    let response = agent
+        .prompt("Which rust example is best suited for the operation 1 + 2")
+        .await?;
+
+    println!("{}", response);
+
+    Ok(())
+}
+
+async fn context() -> Result<(), anyhow::Error> {
+    let model = client().completion_model(providers::xai::GROK_BETA);
+
+    // Create an agent with multiple context documents
+    let agent = AgentBuilder::new(model)
+        .context("Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets")
+        .context("Definition of a *glarb-glarb*: A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.")
+        .context("Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.")
+        .build();
+
+    // Prompt the agent and print the response
+    let response = agent.prompt("What does \"glarb-glarb\" mean?").await?;
+
+    println!("{}", response);
+
+    Ok(())
+}
+
+#[derive(Deserialize)]
+struct OperationArgs {
+    x: i32,
+    y: i32,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("Math error")]
+struct MathError;
+
+#[derive(Deserialize, Serialize)]
+struct Adder;
+impl Tool for Adder {
+    const NAME: &'static str = "add";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        ToolDefinition {
+            name: "add".to_string(),
+            description: "Add x and y together".to_string(),
+            parameters: json!({
+                "type": "object",
+                "properties": {
+                    "x": {
+                        "type": "number",
+                        "description": "The first number to add"
+                    },
+                    "y": {
+                        "type": "number",
+                        "description": "The second number to add"
+                    }
+                }
+            }),
+        }
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x + args.y;
+        Ok(result)
+    }
+}
+
+#[derive(Deserialize, Serialize)]
+struct Subtract;
+impl Tool for Subtract {
+    const NAME: &'static str = "subtract";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        serde_json::from_value(json!({
+            "name": "subtract",
+            "description": "Subtract y from x (i.e.: x - y)",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {
+                        "type": "number",
+                        "description": "The number to substract from"
+                    },
+                    "y": {
+                        "type": "number",
+                        "description": "The number to substract"
+                    }
+                }
+            }
+        }))
+        .expect("Tool Definition")
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x - args.y;
+        Ok(result)
+    }
+}

rig-core/examples/xai_embeddings.rs

@@ -0,0 +1,19 @@
+use rig::providers::xai;
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    // Initialize the xAI client
+    let client = xai::Client::from_env();
+
+    let embeddings = client
+        .embeddings(xai::embedding::EMBEDDING_V1)
+        .simple_document("doc0", "Hello, world!")
+        .simple_document("doc1", "Goodbye, world!")
+        .build()
+        .await
+        .expect("Failed to embed documents");
+
+    println!("{:?}", embeddings);
+
+    Ok(())
+}

rig-core/src/providers/mod.rs

@@ -45,3 +45,4 @@ pub mod cohere;
 pub mod gemini;
 pub mod openai;
 pub mod perplexity;
+pub mod xai;