reagent_rs/services/llm/
client.rs

1use std::{pin::Pin, sync::Arc};
2
3use futures::Stream;
4
5use crate::{
6    services::llm::{
7        models::{
8            chat::{ChatRequest, ChatResponse, ChatStreamChunk},
9            embedding::{EmbeddingsRequest, EmbeddingsResponse},
10            errors::InferenceClientError,
11        },
12        SchemaSpec, StructuredOuputFormat,
13    },
14    ClientConfig,
15};
16
17use super::providers::{
18    anthropic::AnthropicClient, mistral::MistralClient, ollama::OllamaClient, openai::OpenAiClient,
19    openrouter::OpenRouterClient,
20};
21
22#[derive(Debug, Clone, Default)]
23pub enum Provider {
24    #[default]
25    Ollama,
26    OpenAi,
27    Mistral,
28    Anthropic,
29    OpenRouter,
30}
31
32#[derive(Debug, Clone)]
33enum ClientInner {
34    Ollama(OllamaClient),
35    OpenAi(OpenAiClient),
36    Mistral(MistralClient),
37    Anthropic(AnthropicClient),
38    OpenRouter(OpenRouterClient),
39}
40
41#[derive(Clone, Debug)]
42pub struct InferenceClient {
43    config: ClientConfig,
44    inner: Arc<ClientInner>,
45}
46
47impl InferenceClient {
48    pub fn get_config(&self) -> &ClientConfig {
49        &self.config
50    }
51
52    pub fn structured_output_format(
53        &self,
54        spec: &SchemaSpec,
55    ) -> Result<serde_json::Value, InferenceClientError> {
56        match self.get_config().provider {
57            Some(Provider::Ollama) => Ok(OllamaClient::format(spec)),
58            Some(Provider::OpenAi) => Ok(OpenAiClient::format(spec)),
59            Some(Provider::OpenRouter) => Ok(OpenRouterClient::format(spec)),
60            _ => Err(InferenceClientError::Unsupported(
61                "Structured outputs not yet supported for this provider".into(),
62            )),
63        }
64    }
65
66    pub async fn chat(&self, req: ChatRequest) -> Result<ChatResponse, InferenceClientError> {
67        match &*self.inner {
68            ClientInner::Ollama(c) => c.chat(req).await,
69            ClientInner::OpenAi(c) => c.chat(req).await,
70            ClientInner::Mistral(c) => c.chat(req).await,
71            ClientInner::Anthropic(c) => c.chat(req).await,
72            ClientInner::OpenRouter(c) => c.chat(req).await,
73        }
74    }
75
76    pub async fn chat_stream(
77        &self,
78        req: ChatRequest,
79    ) -> Result<
80        Pin<Box<dyn Stream<Item = Result<ChatStreamChunk, InferenceClientError>> + Send + 'static>>,
81        InferenceClientError,
82    > {
83        match &*self.inner {
84            ClientInner::Ollama(c) => c.chat_stream(req).await,
85            ClientInner::OpenAi(c) => c.chat_stream(req).await,
86            ClientInner::Mistral(c) => c.chat_stream(req).await,
87            ClientInner::Anthropic(c) => c.chat_stream(req).await,
88            ClientInner::OpenRouter(c) => c.chat_stream(req).await,
89        }
90    }
91
92    pub async fn embeddings(
93        &self,
94        req: EmbeddingsRequest,
95    ) -> Result<EmbeddingsResponse, InferenceClientError> {
96        match &*self.inner {
97            ClientInner::Ollama(c) => c.embeddings(req).await,
98            ClientInner::OpenAi(c) => c.embeddings(req).await,
99            ClientInner::Mistral(c) => c.embeddings(req).await,
100            ClientInner::Anthropic(c) => c.embeddings(req).await,
101            ClientInner::OpenRouter(c) => c.embeddings(req).await,
102        }
103    }
104}
105
106impl TryFrom<ClientConfig> for InferenceClient {
107    type Error = InferenceClientError;
108
109    fn try_from(cfg: ClientConfig) -> Result<Self, Self::Error> {
110        let config = cfg.clone();
111        let Some(provider) = cfg.provider.clone() else {
112            return Err(InferenceClientError::Config("Provider not defined".into()));
113        };
114        let inner = match provider {
115            Provider::Ollama => ClientInner::Ollama(OllamaClient::new(cfg)?),
116            Provider::OpenAi => ClientInner::OpenAi(OpenAiClient::new(cfg)?),
117            Provider::Mistral => ClientInner::Mistral(MistralClient::new(cfg)?),
118            Provider::Anthropic => ClientInner::Anthropic(AnthropicClient::new(cfg)?),
119            Provider::OpenRouter => ClientInner::OpenRouter(OpenRouterClient::new(cfg)?),
120        };
121        Ok(Self {
122            config,
123            inner: Arc::new(inner),
124        })
125    }
126}
reagent_rs/services/llm/client.rs

reagent_rs/services/llm/
client.rs
