reagent-rs 0.2.9

A Rust library for building AI agents with MCP, custom tools and skills
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

use std::{pin::Pin, sync::Arc};

use futures::Stream;

use crate::{
    services::llm::{
        models::{
            chat::{ChatRequest, ChatResponse, ChatStreamChunk},
            embedding::{EmbeddingsRequest, EmbeddingsResponse},
            errors::InferenceClientError,
        },
        SchemaSpec, StructuredOuputFormat,
    },
    ClientConfig,
};

use super::providers::{
    anthropic::AnthropicClient, mistral::MistralClient, ollama::OllamaClient, openai::OpenAiClient,
    openrouter::OpenRouterClient,
};

#[derive(Debug, Clone, Default)]
pub enum Provider {
    #[default]
    Ollama,
    OpenAi,
    Mistral,
    Anthropic,
    OpenRouter,
}

#[derive(Debug, Clone)]
enum ClientInner {
    Ollama(OllamaClient),
    OpenAi(OpenAiClient),
    Mistral(MistralClient),
    Anthropic(AnthropicClient),
    OpenRouter(OpenRouterClient),
}

#[derive(Clone, Debug)]
pub struct InferenceClient {
    config: ClientConfig,
    inner: Arc<ClientInner>,
}

impl InferenceClient {
    pub fn get_config(&self) -> &ClientConfig {
        &self.config
    }

    pub fn structured_output_format(
        &self,
        spec: &SchemaSpec,
    ) -> Result<serde_json::Value, InferenceClientError> {
        match self.get_config().provider {
            Some(Provider::Ollama) => Ok(OllamaClient::format(spec)),
            Some(Provider::OpenAi) => Ok(OpenAiClient::format(spec)),
            Some(Provider::OpenRouter) => Ok(OpenRouterClient::format(spec)),
            _ => Err(InferenceClientError::Unsupported(
                "Structured outputs not yet supported for this provider".into(),
            )),
        }
    }

    pub async fn chat(&self, req: ChatRequest) -> Result<ChatResponse, InferenceClientError> {
        match &*self.inner {
            ClientInner::Ollama(c) => c.chat(req).await,
            ClientInner::OpenAi(c) => c.chat(req).await,
            ClientInner::Mistral(c) => c.chat(req).await,
            ClientInner::Anthropic(c) => c.chat(req).await,
            ClientInner::OpenRouter(c) => c.chat(req).await,
        }
    }

    pub async fn chat_stream(
        &self,
        req: ChatRequest,
    ) -> Result<
        Pin<Box<dyn Stream<Item = Result<ChatStreamChunk, InferenceClientError>> + Send + 'static>>,
        InferenceClientError,
    > {
        match &*self.inner {
            ClientInner::Ollama(c) => c.chat_stream(req).await,
            ClientInner::OpenAi(c) => c.chat_stream(req).await,
            ClientInner::Mistral(c) => c.chat_stream(req).await,
            ClientInner::Anthropic(c) => c.chat_stream(req).await,
            ClientInner::OpenRouter(c) => c.chat_stream(req).await,
        }
    }

    pub async fn embeddings(
        &self,
        req: EmbeddingsRequest,
    ) -> Result<EmbeddingsResponse, InferenceClientError> {
        match &*self.inner {
            ClientInner::Ollama(c) => c.embeddings(req).await,
            ClientInner::OpenAi(c) => c.embeddings(req).await,
            ClientInner::Mistral(c) => c.embeddings(req).await,
            ClientInner::Anthropic(c) => c.embeddings(req).await,
            ClientInner::OpenRouter(c) => c.embeddings(req).await,
        }
    }
}

impl TryFrom<ClientConfig> for InferenceClient {
    type Error = InferenceClientError;

    fn try_from(cfg: ClientConfig) -> Result<Self, Self::Error> {
        let config = cfg.clone();
        let Some(provider) = cfg.provider.clone() else {
            return Err(InferenceClientError::Config("Provider not defined".into()));
        };
        let inner = match provider {
            Provider::Ollama => ClientInner::Ollama(OllamaClient::new(cfg)?),
            Provider::OpenAi => ClientInner::OpenAi(OpenAiClient::new(cfg)?),
            Provider::Mistral => ClientInner::Mistral(MistralClient::new(cfg)?),
            Provider::Anthropic => ClientInner::Anthropic(AnthropicClient::new(cfg)?),
            Provider::OpenRouter => ClientInner::OpenRouter(OpenRouterClient::new(cfg)?),
        };
        Ok(Self {
            config,
            inner: Arc::new(inner),
        })
    }
}