rig-core 0.10.0

An opinionated library for building LLM powered applications.
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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use rig::{
    agent::AgentBuilder,
    completion::{Prompt, ToolDefinition},
    providers::together,
    tool::Tool,
};
use serde::{Deserialize, Serialize};
use serde_json::json;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    println!("Running basic agent with together");
    basic().await?;

    println!("\nRunning tools agent with tools");
    tools().await?;

    println!("\nRunning together agent with context");
    context().await?;

    println!("\n\nAll agents ran successfully");
    Ok(())
}

async fn basic() -> Result<(), anyhow::Error> {
    let together_ai_client = together::Client::new(
        &std::env::var("TOGETHER_API_KEY").expect("TOGETHER_API_KEY not set"),
    );

    // Choose a model, replace "together-model-v1" with an actual Together AI model name
    let model =
        together_ai_client.completion_model(rig::providers::together::MIXTRAL_8X7B_INSTRUCT_V0_1);

    let agent = AgentBuilder::new(model)
        .preamble("You are a comedian here to entertain the user using humour and jokes.")
        .build();

    // Prompt the agent and print the response
    let response = agent.prompt("Entertain me!").await?;
    println!("{}", response);

    Ok(())
}

async fn tools() -> Result<(), anyhow::Error> {
    // Create Together AI client
    let together_ai_client = together::Client::new(
        &std::env::var("TOGETHER_API_KEY").expect("TOGETHER_API_KEY not set"),
    );

    // Choose a model, replace "together-model-v1" with an actual Together AI model name
    let model =
        together_ai_client.completion_model(rig::providers::together::MIXTRAL_8X7B_INSTRUCT_V0_1);

    // Create an agent with multiple context documents
    let calculator_agent = AgentBuilder::new(model)
        .preamble("You are a calculator here to help the user perform arithmetic operations. Use the tools provided to answer the user's question.")
        .tool(Adder)
        .build();

    // Prompt the agent and print the response
    println!("Calculate 5 + 3");
    println!(
        "Calculator Agent: {}",
        calculator_agent.prompt("Calculate 5 + 3").await?
    );

    Ok(())
}

async fn context() -> Result<(), anyhow::Error> {
    // Create Together AI client
    let together_ai_client = together::Client::new(
        &std::env::var("TOGETHER_API_KEY").expect("TOGETHER_API_KEY not set"),
    );

    // Choose a model, replace "together-model-v1" with an actual Together AI model name
    let model =
        together_ai_client.completion_model(rig::providers::together::MIXTRAL_8X7B_INSTRUCT_V0_1);

    // 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 an 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(Debug, 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> {
        println!("The args: {:?}", args);
        let result = args.x + args.y;
        Ok(result)
    }
}