rigs 0.0.8

A orchestration framework for rig
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

use rig::{
    agent::AgentBuilder,
    providers::{
        anthropic,
        deepseek::{self, DeepSeekCompletionModel},
        gemini, openai, openrouter,
    },
};
use thiserror::Error;

#[derive(Clone)]
pub enum LLMProvider {
    Anthropic(ModelConfig),
    DeepSeek(ModelConfig),
    Gemini(ModelConfig),
    OpenAI(ModelConfig),
    OpenRouter(ModelConfig),
}

macro_rules! impl_agent_builder {
    ($method:ident, $variant:ident, $client:ty, $model:ty) => {
        pub fn $method(&self) -> Result<AgentBuilder<$model>, LLMProviderError> {
            let LLMProvider::$variant(config) = self else {
                return Err(LLMProviderError::LLMProviderNotMatch);
            };
            let client = <$client>::from_env();
            Ok(client.agent(&config.model))
        }
    };
}

macro_rules! impl_agent_builder_auto {
    ($variant:ident, $client:ty, $model:ty) => {
        paste::paste! {
            impl_agent_builder!(
                [<get_ $variant:snake _agent_builder>],
                $variant,
                $client,
                $model
            );
        }
    };
}

impl LLMProvider {
    pub fn anthropic(model: impl Into<String>) -> Self {
        Self::Anthropic(ModelConfig {
            model: model.into(),
            stream: false,
        })
    }

    pub fn deepseek(model: impl Into<String>) -> Self {
        Self::DeepSeek(ModelConfig {
            model: model.into(),
            stream: false,
        })
    }

    pub fn gemini(model: impl Into<String>) -> Self {
        Self::Gemini(ModelConfig {
            model: model.into(),
            stream: false,
        })
    }

    pub fn openai(model: impl Into<String>) -> Self {
        Self::OpenAI(ModelConfig {
            model: model.into(),
            stream: false,
        })
    }

    pub fn openrouter(model: impl Into<String>) -> Self {
        Self::OpenRouter(ModelConfig {
            model: model.into(),
            stream: false,
        })
    }

    pub fn get_config(&self) -> &ModelConfig {
        match self {
            LLMProvider::Anthropic(config)
            | LLMProvider::DeepSeek(config)
            | LLMProvider::Gemini(config)
            | LLMProvider::OpenAI(config)
            | LLMProvider::OpenRouter(config) => config,
        }
    }

    impl_agent_builder_auto!(
        Anthropic,
        anthropic::Client,
        anthropic::completion::CompletionModel
    );

    impl_agent_builder_auto!(DeepSeek, deepseek::Client, DeepSeekCompletionModel);

    impl_agent_builder_auto!(Gemini, gemini::Client, gemini::completion::CompletionModel);

    impl_agent_builder_auto!(OpenAI, openai::Client, openai::CompletionModel);

    impl_agent_builder_auto!(OpenRouter, openrouter::Client, openrouter::CompletionModel);
}

#[derive(Clone)]
pub struct ModelConfig {
    pub model: String,
    pub stream: bool,
}

#[derive(Debug, Error)]
pub enum LLMProviderError {
    #[error("LLM provider not match")]
    LLMProviderNotMatch,
}