mentra/
provider.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5
6pub use mentra_provider::AnthropicRequestOptions;
7pub use mentra_provider::AuthScheme;
8pub use mentra_provider::BuiltinProvider;
9pub use mentra_provider::CompactionInputItem;
10pub use mentra_provider::CompactionRequest;
11pub use mentra_provider::CompactionResponse;
12pub use mentra_provider::ContentBlock;
13pub use mentra_provider::ContentBlockDelta;
14pub use mentra_provider::ContentBlockStart;
15pub use mentra_provider::EmbeddingData;
16pub use mentra_provider::EmbeddingModelInfo;
17pub use mentra_provider::EmbeddingProvider;
18pub use mentra_provider::EmbeddingRequest;
19pub use mentra_provider::EmbeddingResponse;
20pub use mentra_provider::EmbeddingUsage;
21pub use mentra_provider::GeminiRequestOptions;
22pub use mentra_provider::ImageSource;
23pub use mentra_provider::MemorySummarizeOutput;
24pub use mentra_provider::MemorySummarizeRequest;
25pub use mentra_provider::MemorySummarizeResponse;
26pub use mentra_provider::Message;
27pub use mentra_provider::ModelInfo;
28pub use mentra_provider::ModelSelector;
29pub use mentra_provider::OpenAIRequestOptions;
30pub use mentra_provider::ProviderCapabilities;
31pub use mentra_provider::ProviderCredentials;
32pub use mentra_provider::ProviderDefinition;
33pub use mentra_provider::ProviderDescriptor;
34pub use mentra_provider::ProviderError;
35pub use mentra_provider::ProviderEvent;
36pub use mentra_provider::ProviderEventStream;
37pub use mentra_provider::ProviderId;
38pub use mentra_provider::ProviderRequestOptions;
39pub use mentra_provider::RawMemory;
40pub use mentra_provider::RawMemoryMetadata;
41pub use mentra_provider::ReasoningEffort;
42pub use mentra_provider::ReasoningFormat;
43pub use mentra_provider::ReasoningOptions;
44pub use mentra_provider::ReasoningProvenance;
45pub use mentra_provider::Request;
46pub use mentra_provider::RequestFields;
47pub use mentra_provider::Response;
48pub use mentra_provider::ResponsesRequestOptions;
49pub use mentra_provider::ResponsesStateMode;
50pub use mentra_provider::ResponsesTransport;
51pub use mentra_provider::Role;
52pub use mentra_provider::TokenUsage;
53pub use mentra_provider::ToolChoice;
54pub use mentra_provider::ToolSearchMode;
55pub use mentra_provider::WireApi;
56pub use mentra_provider::collect_response_from_stream;
57pub use mentra_provider::provider_event_stream_from_response;
58
59/// Transport-neutral interface implemented by model providers.
60#[async_trait]
61pub trait Provider: Send + Sync {
62    /// Returns identifying metadata for the provider instance.
63    fn descriptor(&self) -> ProviderDescriptor;
64
65    /// Returns feature flags supported by this provider instance.
66    fn capabilities(&self) -> ProviderCapabilities {
67        ProviderCapabilities::default()
68    }
69
70    /// Returns the same configured provider with an independent session scope.
71    ///
72    /// Scope creation is synchronous and local: it must not connect to the
73    /// provider or perform other network I/O. Ordinary clones of the returned
74    /// value share that one scope; calling this method again creates an
75    /// independent one with the same descriptor identity.
76    ///
77    /// Custom providers remain valid for one-shot runtimes without implementing
78    /// this method. They report the capability as unsupported until they can
79    /// separate reusable configuration from provider-owned session state.
80    fn fresh_session_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
81        Err(ProviderError::UnsupportedCapability(
82            "fresh_session_scope".to_string(),
83        ))
84    }
85
86    /// Returns the same configured provider with an independent conversation,
87    /// still sharing this one's transport.
88    ///
89    /// The narrower half of [`fresh_session_scope`](Self::fresh_session_scope),
90    /// under the same rules. It separates what answers questions about one
91    /// exchange — where its response chain is, which turn the endpoint routed
92    /// it to — and keeps what describes the endpoint, including any cached
93    /// connection. The runtime mints one of these per agent, so two agents
94    /// never read each other's chain and still share one warm socket.
95    ///
96    /// Custom providers remain valid without implementing this method; they
97    /// report the capability as unsupported and the runtime keeps the scope it
98    /// already had.
99    fn fresh_conversation_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
100        Err(ProviderError::UnsupportedCapability(
101            "fresh_conversation_scope".to_string(),
102        ))
103    }
104
105    /// Lists models available from the provider.
106    async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError>;
107
108    /// Streams a model response for the given request.
109    async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError>;
110
111    /// Sends a request and collects the full response in memory.
112    async fn send(&self, request: Request<'_>) -> Result<Response, ProviderError> {
113        collect_response_from_stream(self.stream(request).await?).await
114    }
115
116    /// Compacts transcript history using a provider-native endpoint when supported.
117    async fn compact(
118        &self,
119        _request: CompactionRequest<'_>,
120    ) -> Result<CompactionResponse, ProviderError> {
121        Err(ProviderError::UnsupportedCapability(
122            "history_compaction".to_string(),
123        ))
124    }
125
126    /// Summarizes raw trace memories using a provider-native implementation when supported.
127    async fn summarize_memories(
128        &self,
129        _request: MemorySummarizeRequest<'_>,
130    ) -> Result<MemorySummarizeResponse, ProviderError> {
131        Err(ProviderError::UnsupportedCapability(
132            "memory_summarization".to_string(),
133        ))
134    }
135}
136
137/// One high-level provider bound to one provider-owned session-state scope.
138///
139/// Cloning this value deliberately shares the current scope. Use
140/// [`Provider::fresh_session_scope`] to retain provider configuration while
141/// allocating independent turn, response-chain, connection, and in-flight
142/// state. Neither operation implicitly opens or warms a connection.
143#[derive(Clone)]
144pub struct ProviderSessionScope {
145    inner: Arc<dyn Provider>,
146}
147
148impl ProviderSessionScope {
149    /// Wraps a provider whose current session state is the scope to share.
150    pub fn new<P>(provider: P) -> Self
151    where
152        P: Provider + 'static,
153    {
154        Self {
155            inner: Arc::new(provider),
156        }
157    }
158}
159
160#[async_trait]
161impl Provider for ProviderSessionScope {
162    fn descriptor(&self) -> ProviderDescriptor {
163        self.inner.descriptor()
164    }
165
166    fn capabilities(&self) -> ProviderCapabilities {
167        self.inner.capabilities()
168    }
169
170    fn fresh_session_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
171        self.inner.fresh_session_scope()
172    }
173
174    fn fresh_conversation_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
175        self.inner.fresh_conversation_scope()
176    }
177
178    async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
179        self.inner.list_models().await
180    }
181
182    async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
183        self.inner.stream(request).await
184    }
185
186    async fn send(&self, request: Request<'_>) -> Result<Response, ProviderError> {
187        self.inner.send(request).await
188    }
189
190    async fn compact(
191        &self,
192        request: CompactionRequest<'_>,
193    ) -> Result<CompactionResponse, ProviderError> {
194        self.inner.compact(request).await
195    }
196
197    async fn summarize_memories(
198        &self,
199        request: MemorySummarizeRequest<'_>,
200    ) -> Result<MemorySummarizeResponse, ProviderError> {
201        self.inner.summarize_memories(request).await
202    }
203}
204
205/// Forwards to the provider behind the pointer, so an `Arc<dyn Provider>` —
206/// what the builtin provider constructors hand back — can be passed straight to
207/// the generic `P: Provider` entry points.
208#[async_trait]
209impl Provider for Arc<dyn Provider> {
210    fn descriptor(&self) -> ProviderDescriptor {
211        (**self).descriptor()
212    }
213
214    fn capabilities(&self) -> ProviderCapabilities {
215        (**self).capabilities()
216    }
217
218    fn fresh_session_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
219        (**self).fresh_session_scope()
220    }
221
222    fn fresh_conversation_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
223        (**self).fresh_conversation_scope()
224    }
225
226    async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
227        (**self).list_models().await
228    }
229
230    async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
231        (**self).stream(request).await
232    }
233
234    async fn send(&self, request: Request<'_>) -> Result<Response, ProviderError> {
235        (**self).send(request).await
236    }
237
238    async fn compact(
239        &self,
240        request: CompactionRequest<'_>,
241    ) -> Result<CompactionResponse, ProviderError> {
242        (**self).compact(request).await
243    }
244
245    async fn summarize_memories(
246        &self,
247        request: MemorySummarizeRequest<'_>,
248    ) -> Result<MemorySummarizeResponse, ProviderError> {
249        (**self).summarize_memories(request).await
250    }
251}
252
253#[derive(Default)]
254pub struct ProviderRegistry {
255    default_provider: Option<ProviderId>,
256    default_embedding_provider: Option<ProviderId>,
257    providers: HashMap<ProviderId, Arc<dyn Provider>>,
258    embedding_providers: HashMap<ProviderId, Arc<dyn EmbeddingProvider>>,
259    /// The Responses transport this runtime's requests go out on, or `None`
260    /// when the runtime does not choose and each request's own options stand.
261    ///
262    /// It lives here rather than on the handle because a transport is a
263    /// property of the connection to a provider, and this is where a runtime
264    /// keeps those. It also means the choice travels the one path the builder
265    /// already hands to the handle at build time, instead of a field every
266    /// `with_*` reconstructor would have to remember to carry.
267    responses_transport: Option<ResponsesTransport>,
268    responses_state_mode: Option<ResponsesStateMode>,
269}
270
271impl ProviderRegistry {
272    pub(crate) fn register_builtin_provider(
273        &mut self,
274        id: BuiltinProvider,
275        api_key: impl Into<String>,
276    ) -> Result<(), String> {
277        let api_key = api_key.into();
278        let provider: Arc<dyn Provider> = match id {
279            BuiltinProvider::Anthropic => anthropic::provider(api_key.clone()),
280            BuiltinProvider::Gemini => gemini::provider(api_key.clone()),
281            BuiltinProvider::OpenAI => openai::provider(api_key.clone()),
282            BuiltinProvider::OpenRouter => openrouter::provider(api_key.clone()),
283            BuiltinProvider::Ollama => ollama::provider(),
284            BuiltinProvider::LmStudio => lmstudio::provider(),
285        };
286
287        let provider_id: ProviderId = id.into();
288
289        if self.default_provider.is_none() {
290            self.default_provider = Some(provider_id.clone());
291        }
292
293        // Register embedding provider for providers that support it.
294        let ep: Option<Arc<dyn EmbeddingProvider>> = match id {
295            BuiltinProvider::OpenAI => Some(Arc::new(mentra_provider::responses::openai(api_key))),
296            BuiltinProvider::OpenRouter => {
297                Some(Arc::new(mentra_provider::responses::openrouter(api_key)))
298            }
299            BuiltinProvider::Ollama => Some(Arc::new(openai_compatible_embedding_provider(
300                id,
301                "http://127.0.0.1:11434/",
302            ))),
303            BuiltinProvider::LmStudio => Some(Arc::new(openai_compatible_embedding_provider(
304                id,
305                "http://127.0.0.1:1234/",
306            ))),
307            _ => None,
308        };
309        if let Some(ep) = ep {
310            if self.default_embedding_provider.is_none() {
311                self.default_embedding_provider = Some(provider_id.clone());
312            }
313            self.embedding_providers.insert(provider_id.clone(), ep);
314        }
315
316        self.providers.insert(provider_id, provider);
317        Ok(())
318    }
319
320    pub(crate) fn register_provider_instance<P>(&mut self, provider: P)
321    where
322        P: Provider + 'static,
323    {
324        let descriptor = provider.descriptor();
325        let id = descriptor.id;
326
327        if self.default_provider.is_none() {
328            self.default_provider = Some(id.clone());
329        }
330
331        self.providers.insert(id, Arc::new(provider));
332    }
333
334    pub(crate) fn register_registered_provider<P>(&mut self, provider: P)
335    where
336        P: mentra_provider::Provider + 'static,
337    {
338        let descriptor = provider.descriptor();
339        let id = descriptor.id;
340
341        if self.default_provider.is_none() {
342            self.default_provider = Some(id.clone());
343        }
344
345        self.providers.insert(id, shared_provider(provider));
346    }
347
348    /// Registers an already-shared provider, avoiding a second `Arc` around the
349    /// one the builtin constructors return.
350    pub(crate) fn register_shared_provider(&mut self, provider: Arc<dyn Provider>) {
351        let id = provider.descriptor().id;
352
353        if self.default_provider.is_none() {
354            self.default_provider = Some(id.clone());
355        }
356
357        self.providers.insert(id, provider);
358    }
359
360    pub(crate) fn register_ollama(&mut self) {
361        self.register_shared_provider(ollama::provider());
362    }
363
364    pub(crate) fn register_lmstudio(&mut self) {
365        self.register_shared_provider(lmstudio::provider());
366    }
367
368    pub(crate) fn get_provider(&self, id: Option<&ProviderId>) -> Option<Arc<dyn Provider>> {
369        match id {
370            Some(id) => self.providers.get(id).cloned(),
371            None => self
372                .default_provider
373                .as_ref()
374                .and_then(|id| self.providers.get(id).cloned()),
375        }
376    }
377
378    /// Returns the default embedding provider, or `None` if no embedding-capable provider
379    /// has been registered.
380    ///
381    /// The default is the first embedding-capable provider registered. To look up a
382    /// specific provider use [`embedding_provider_for`](Self::embedding_provider_for).
383    pub fn embedding_provider(&self) -> Option<Arc<dyn EmbeddingProvider>> {
384        self.default_embedding_provider
385            .as_ref()
386            .and_then(|id| self.embedding_providers.get(id))
387            .map(Arc::clone)
388            .or_else(|| self.embedding_providers.values().next().map(Arc::clone))
389    }
390
391    /// Returns the embedding provider for a specific provider ID, or `None`.
392    pub fn embedding_provider_for(&self, id: &ProviderId) -> Option<Arc<dyn EmbeddingProvider>> {
393        self.embedding_providers.get(id).map(Arc::clone)
394    }
395
396    pub(crate) fn descriptors(&self) -> Vec<ProviderDescriptor> {
397        self.providers
398            .values()
399            .map(|provider| provider.descriptor())
400            .collect()
401    }
402
403    pub(crate) fn is_empty(&self) -> bool {
404        self.providers.is_empty()
405    }
406
407    pub(crate) fn set_responses_transport(&mut self, transport: ResponsesTransport) {
408        self.responses_transport = Some(transport);
409    }
410
411    pub(crate) fn responses_transport(&self) -> Option<ResponsesTransport> {
412        self.responses_transport
413    }
414
415    pub(crate) fn set_responses_state_mode(&mut self, state_mode: ResponsesStateMode) {
416        self.responses_state_mode = Some(state_mode);
417    }
418
419    pub(crate) fn responses_state_mode(&self) -> Option<ResponsesStateMode> {
420        self.responses_state_mode
421    }
422}
423
424/// Settles which Responses transport a request goes out on, and refuses one the
425/// provider cannot serve.
426///
427/// The runtime's choice, when it made one, replaces whatever the request's own
428/// options carried: it is the connection-level answer, and a per-request one
429/// that disagreed would mean two live opinions about a single socket. With no
430/// runtime choice the request's own value stands, which is what every caller
431/// had before a runtime could choose at all.
432///
433/// A provider whose capabilities report no websocket support is refused rather
434/// than quietly served over HTTP+SSE. The fallback is the tempting behavior and
435/// the wrong one: asking for a transport is explicit, so answering on a
436/// different one returns a stream nobody asked for and hides a misconfigured
437/// runtime behind a working one — the same stance `stream_response` already
438/// takes when the transport is not compiled in.
439/// Gives one agent its own conversation on the provider it was handed.
440///
441/// Every agent owns a transcript, so every agent owns the state that answers
442/// questions about that transcript: where its response chain is, and which
443/// turn the endpoint routed it to. A registry hands out one `Arc` per
444/// provider, so without this two agents built from the same registry entry
445/// read and overwrite each other's — and a host that runs a conversation per
446/// chat has each chat chaining from whichever one answered last.
447///
448/// The transport is deliberately not split. A cached websocket describes the
449/// endpoint rather than the exchange, and an agent-per-turn host would
450/// otherwise pay a handshake per turn; [`Provider::fresh_session_scope`]
451/// remains the way to separate that too.
452///
453/// A provider that declines keeps the scope it was handed. Declining is not a
454/// failure to report: minting is defined as local and synchronous, so there is
455/// no transport error to surface, and the one defined refusal —
456/// `UnsupportedCapability` — says the provider holds no per-conversation state,
457/// which is true of every provider that does not chain.
458pub(crate) fn conversation_scoped(provider: Arc<dyn Provider>) -> Arc<dyn Provider> {
459    match provider.fresh_conversation_scope() {
460        Ok(scope) => Arc::new(scope),
461        Err(_) => provider,
462    }
463}
464
465/// Applies the runtime's state-mode choice, when it made one.
466///
467/// Same rule as [`select_responses_transport`]: a runtime-level choice is the
468/// connection-level answer and replaces whatever the request's own options
469/// carried, and with no runtime choice the request's own value stands — which
470/// is [`ResponsesStateMode::ReplayOnly`] unless an agent's
471/// `provider_request_options` said otherwise.
472///
473/// Unlike a transport, nothing here can be refused: every Responses endpoint
474/// accepts a request that chains nothing, and a `previous_response_id` the
475/// endpoint will not take is already recovered at the session.
476pub(crate) fn select_responses_state_mode(
477    chosen: Option<ResponsesStateMode>,
478    options: &mut ProviderRequestOptions,
479) {
480    if let Some(state_mode) = chosen {
481        options.responses.state_mode = state_mode;
482    }
483}
484
485pub(crate) fn select_responses_transport(
486    provider: &dyn Provider,
487    chosen: Option<ResponsesTransport>,
488    options: &mut ProviderRequestOptions,
489) -> Result<(), crate::error::RuntimeError> {
490    if let Some(transport) = chosen {
491        options.responses.transport = transport;
492    }
493
494    if options.responses.transport != ResponsesTransport::WebSocket
495        || provider.capabilities().supports_websockets
496    {
497        return Ok(());
498    }
499
500    let descriptor = provider.descriptor();
501    let name = descriptor
502        .display_name
503        .unwrap_or_else(|| descriptor.id.as_str().to_string());
504    Err(crate::error::RuntimeError::OperationDenied(format!(
505        "provider '{name}' does not serve the Responses websocket transport; \
506         select ResponsesTransport::HttpSse or register a provider that does \
507         — answering over HTTP+SSE would return a transport nobody asked for"
508    )))
509}
510
511fn shared_provider<P>(provider: P) -> Arc<dyn Provider>
512where
513    P: mentra_provider::Provider + 'static,
514{
515    Arc::new(SharedProviderProxy { inner: provider })
516}
517
518/// Builds a `ResponsesProvider` (with no credentials) for OpenAI-compatible
519/// local providers (Ollama, LmStudio) so they can be used as embedding providers.
520fn openai_compatible_embedding_provider(
521    builtin: BuiltinProvider,
522    base_url: &str,
523) -> mentra_provider::responses::ResponsesProvider<NoCredentialsSource> {
524    use mentra_provider::AuthScheme;
525    use mentra_provider::ProviderCapabilities;
526    use mentra_provider::WireApi;
527    use std::collections::HashMap;
528
529    let mut definition = ProviderDefinition::new(builtin);
530    definition.wire_api = WireApi::Responses;
531    definition.auth_scheme = AuthScheme::None;
532    definition.capabilities = ProviderCapabilities {
533        supports_model_listing: true,
534        supports_streaming: true,
535        supports_websockets: false,
536        supports_tool_calls: true,
537        supports_images: true,
538        supports_history_compaction: false,
539        supports_memory_summarization: false,
540        supports_deferred_tools: false,
541        supports_hosted_tool_search: false,
542        supports_hosted_web_search: false,
543        supports_image_generation: false,
544        supports_reasoning_effort: false,
545        reports_reasoning_tokens: false,
546        reports_thoughts_tokens: false,
547        supports_structured_tool_results: false,
548        supports_embeddings: true,
549    };
550    definition.base_url = Some(base_url.to_string());
551    definition.headers = Some(HashMap::new());
552    mentra_provider::responses::ResponsesProvider::new(definition, NoCredentialsSource)
553}
554
555/// Builds a provider for an endpoint speaking the OpenAI `chat/completions`
556/// wire.
557///
558/// This is the wire the OpenAI-compatible ecosystem actually implements.
559/// `v1/responses` is OpenAI's own, and an endpoint that does not serve it
560/// answers 404 to every turn — with an error that reads like a mistyped base
561/// URL rather than like a wire mismatch.
562fn chat_completions_provider(
563    provider: impl Into<mentra_provider::ProviderId>,
564    display_name: &str,
565    description: &str,
566    base_url: &str,
567    credentials: Option<String>,
568) -> Arc<dyn Provider> {
569    let mut definition = mentra_provider::chat_completions::definition(provider, base_url);
570    definition.descriptor.display_name = Some(display_name.to_string());
571    definition.descriptor.description = Some(description.to_string());
572
573    match credentials {
574        Some(api_key) => shared_provider(
575            mentra_provider::chat_completions::ChatCompletionsProvider::new(
576                definition,
577                mentra_provider::StaticCredentialSource::new(api_key),
578            ),
579        ),
580        None => {
581            definition.auth_scheme = mentra_provider::AuthScheme::None;
582            shared_provider(
583                mentra_provider::chat_completions::ChatCompletionsProvider::new(
584                    definition,
585                    NoCredentialsSource,
586                ),
587            )
588        }
589    }
590}
591
592#[derive(Clone)]
593struct NoCredentialsSource;
594
595#[async_trait]
596impl mentra_provider::CredentialSource for NoCredentialsSource {
597    async fn credentials(
598        &self,
599    ) -> Result<mentra_provider::ProviderCredentials, mentra_provider::ProviderError> {
600        Ok(mentra_provider::ProviderCredentials::default())
601    }
602}
603
604struct SharedProviderProxy<P> {
605    inner: P,
606}
607
608#[async_trait]
609impl<P> Provider for SharedProviderProxy<P>
610where
611    P: mentra_provider::Provider + 'static,
612{
613    fn descriptor(&self) -> ProviderDescriptor {
614        self.inner.descriptor()
615    }
616
617    fn capabilities(&self) -> ProviderCapabilities {
618        self.inner.definition().capabilities
619    }
620
621    fn fresh_session_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
622        let scope = mentra_provider::Provider::fresh_session_scope(&self.inner)?;
623        Ok(ProviderSessionScope::new(SharedProviderProxy {
624            inner: scope,
625        }))
626    }
627
628    fn fresh_conversation_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
629        let scope = mentra_provider::Provider::fresh_conversation_scope(&self.inner)?;
630        Ok(ProviderSessionScope::new(SharedProviderProxy {
631            inner: scope,
632        }))
633    }
634
635    async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
636        self.inner.list_models().await
637    }
638
639    async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
640        self.inner.stream(request).await
641    }
642
643    async fn compact(
644        &self,
645        request: CompactionRequest<'_>,
646    ) -> Result<CompactionResponse, ProviderError> {
647        self.inner.compact(request).await
648    }
649
650    async fn summarize_memories(
651        &self,
652        request: MemorySummarizeRequest<'_>,
653    ) -> Result<MemorySummarizeResponse, ProviderError> {
654        self.inner.summarize_memories(request).await
655    }
656}
657
658pub mod openai {
659    use std::sync::Arc;
660
661    use async_trait::async_trait;
662
663    use super::Provider;
664    use super::shared_provider;
665
666    /// Supplies OpenAI API credentials on demand.
667    #[async_trait]
668    pub trait OpenAICredentialSource: Send + Sync {
669        async fn api_key(&self) -> Result<String, String>;
670    }
671
672    /// Builds the OpenAI provider from a static API key.
673    pub fn provider(api_key: impl Into<String>) -> Arc<dyn Provider> {
674        shared_provider(mentra_provider::responses::openai(api_key))
675    }
676
677    /// Builds the OpenAI provider from a credential source resolved per request.
678    pub fn with_credential_source(
679        source: impl OpenAICredentialSource + 'static,
680    ) -> Arc<dyn Provider> {
681        with_shared_credential_source(Arc::new(source))
682    }
683
684    /// Builds the OpenAI provider from an already-shared credential source.
685    pub fn with_shared_credential_source(
686        source: Arc<dyn OpenAICredentialSource>,
687    ) -> Arc<dyn Provider> {
688        shared_provider(mentra_provider::responses::openai_with_credential_source(
689            OpenAICredentialAdapter { source },
690        ))
691    }
692
693    #[derive(Clone)]
694    struct OpenAICredentialAdapter {
695        source: Arc<dyn OpenAICredentialSource>,
696    }
697
698    #[async_trait]
699    impl mentra_provider::CredentialSource for OpenAICredentialAdapter {
700        async fn credentials(
701            &self,
702        ) -> Result<mentra_provider::ProviderCredentials, mentra_provider::ProviderError> {
703            let api_key = self
704                .source
705                .api_key()
706                .await
707                .map_err(mentra_provider::ProviderError::InvalidRequest)?;
708
709            Ok(mentra_provider::ProviderCredentials {
710                bearer_token: Some(api_key),
711                account_id: None,
712                headers: Default::default(),
713            })
714        }
715    }
716}
717
718pub mod openrouter {
719    use std::sync::Arc;
720
721    use super::Provider;
722    use super::shared_provider;
723
724    /// Builds the OpenRouter provider from a static API key.
725    pub fn provider(api_key: impl Into<String>) -> Arc<dyn Provider> {
726        shared_provider(mentra_provider::responses::openrouter(api_key))
727    }
728}
729
730pub mod anthropic {
731    use std::sync::Arc;
732
733    use super::Provider;
734    use super::shared_provider;
735
736    /// Builds the Anthropic provider from a static API key.
737    pub fn provider(api_key: impl Into<String>) -> Arc<dyn Provider> {
738        shared_provider(mentra_provider::anthropic::AnthropicProvider::new(api_key))
739    }
740}
741
742pub mod gemini {
743    use std::sync::Arc;
744
745    use super::Provider;
746    use super::shared_provider;
747
748    /// Builds the Gemini provider from a static API key.
749    pub fn provider(api_key: impl Into<String>) -> Arc<dyn Provider> {
750        shared_provider(mentra_provider::gemini::GeminiProvider::new(api_key))
751    }
752}
753
754/// Providers for any endpoint speaking the OpenAI `chat/completions` wire.
755///
756/// DeepSeek, Groq, Together, Fireworks, Mistral, xAI, OpenRouter, vLLM,
757/// llama.cpp, Ollama and LM Studio all serve this wire; almost none of them
758/// serve `v1/responses`. Point this at a base URL and it will speak the thing
759/// on the other end.
760pub mod openai_compatible {
761    use std::sync::Arc;
762
763    use super::Provider;
764
765    const DESCRIPTION: &str = "OpenAI-compatible chat/completions provider";
766
767    /// Builds a provider for one OpenAI-compatible endpoint that authenticates
768    /// with a bearer token.
769    ///
770    /// `id` is the name the runtime will know this provider by, and can be
771    /// anything not already registered.
772    ///
773    /// ```rust,no_run
774    /// use mentra::provider::openai_compatible;
775    ///
776    /// let deepseek = openai_compatible::new(
777    ///     "deepseek",
778    ///     "https://api.deepseek.com/",
779    ///     std::env::var("DEEPSEEK_API_KEY").unwrap(),
780    /// );
781    /// ```
782    pub fn new(
783        id: impl Into<mentra_provider::ProviderId>,
784        base_url: impl AsRef<str>,
785        api_key: impl Into<String>,
786    ) -> Arc<dyn Provider> {
787        let id = id.into();
788        let display_name = id.as_str().to_string();
789        super::chat_completions_provider(
790            id,
791            &display_name,
792            DESCRIPTION,
793            base_url.as_ref(),
794            Some(api_key.into()),
795        )
796    }
797
798    /// Builds a provider for an endpoint that wants no credentials — a local
799    /// vLLM or llama.cpp server, say.
800    pub fn without_credentials(
801        id: impl Into<mentra_provider::ProviderId>,
802        base_url: impl AsRef<str>,
803    ) -> Arc<dyn Provider> {
804        let id = id.into();
805        let display_name = id.as_str().to_string();
806        super::chat_completions_provider(id, &display_name, DESCRIPTION, base_url.as_ref(), None)
807    }
808}
809
810pub mod ollama {
811    use std::sync::Arc;
812
813    use super::BuiltinProvider;
814    use super::Provider;
815
816    const DEFAULT_BASE_URL: &str = "http://127.0.0.1:11434/";
817
818    /// Builds the Ollama provider against the default local base URL.
819    pub fn provider() -> Arc<dyn Provider> {
820        with_base_url(DEFAULT_BASE_URL)
821    }
822
823    /// Builds the Ollama provider against a specific base URL.
824    pub fn with_base_url(base_url: impl AsRef<str>) -> Arc<dyn Provider> {
825        // Ollama serves `v1/chat/completions` and has never served
826        // `v1/responses`.
827        super::chat_completions_provider(
828            BuiltinProvider::Ollama,
829            "Ollama",
830            "Ollama OpenAI-compatible chat/completions provider",
831            base_url.as_ref(),
832            None,
833        )
834    }
835}
836
837pub mod lmstudio {
838    use std::sync::Arc;
839
840    use super::BuiltinProvider;
841    use super::Provider;
842
843    const DEFAULT_BASE_URL: &str = "http://127.0.0.1:1234/";
844
845    /// Builds the LM Studio provider against the default local base URL.
846    pub fn provider() -> Arc<dyn Provider> {
847        with_base_url(DEFAULT_BASE_URL)
848    }
849
850    /// Builds the LM Studio provider against a specific base URL.
851    pub fn with_base_url(base_url: impl AsRef<str>) -> Arc<dyn Provider> {
852        // LM Studio's OpenAI-compatible surface is `v1/chat/completions`; only
853        // recent builds serve `v1/responses` at all.
854        super::chat_completions_provider(
855            BuiltinProvider::LmStudio,
856            "LM Studio",
857            "LM Studio OpenAI-compatible chat/completions provider",
858            base_url.as_ref(),
859            None,
860        )
861    }
862}
863
864#[cfg(test)]
865mod conversation_scope_tests {
866    use std::sync::Arc;
867    use std::sync::atomic::{AtomicUsize, Ordering};
868
869    use async_trait::async_trait;
870
871    use super::{
872        BuiltinProvider, ModelInfo, Provider, ProviderDescriptor, ProviderError,
873        ProviderEventStream, ProviderSessionScope, Request, conversation_scoped,
874    };
875
876    /// Mints a scope and counts how many it was asked for, so a call site that
877    /// stopped asking shows up as a count that never moves.
878    struct Minting(Arc<AtomicUsize>);
879
880    #[async_trait]
881    impl Provider for Minting {
882        fn descriptor(&self) -> ProviderDescriptor {
883            ProviderDescriptor::new(BuiltinProvider::OpenAI)
884        }
885
886        fn fresh_conversation_scope(&self) -> Result<ProviderSessionScope, ProviderError> {
887            self.0.fetch_add(1, Ordering::SeqCst);
888            Ok(ProviderSessionScope::new(Minting(Arc::clone(&self.0))))
889        }
890
891        async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
892            Ok(Vec::new())
893        }
894
895        async fn stream(
896            &self,
897            _request: Request<'_>,
898        ) -> Result<ProviderEventStream, ProviderError> {
899            unreachable!("no turn is run in these tests")
900        }
901    }
902
903    /// The default: holds no per-conversation state and says so.
904    struct Declining;
905
906    #[async_trait]
907    impl Provider for Declining {
908        fn descriptor(&self) -> ProviderDescriptor {
909            ProviderDescriptor::new(BuiltinProvider::OpenAI)
910        }
911
912        async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
913            Ok(Vec::new())
914        }
915
916        async fn stream(
917            &self,
918            _request: Request<'_>,
919        ) -> Result<ProviderEventStream, ProviderError> {
920            unreachable!("no turn is run in these tests")
921        }
922    }
923
924    #[test]
925    fn a_provider_that_mints_gets_a_new_scope_each_time() {
926        let minted = Arc::new(AtomicUsize::new(0));
927        let shared: Arc<dyn Provider> = Arc::new(Minting(Arc::clone(&minted)));
928
929        let first = conversation_scoped(Arc::clone(&shared));
930        let second = conversation_scoped(Arc::clone(&shared));
931
932        assert_eq!(minted.load(Ordering::SeqCst), 2);
933        assert!(!Arc::ptr_eq(&first, &shared));
934        assert!(!Arc::ptr_eq(&first, &second));
935    }
936
937    #[test]
938    fn a_provider_that_declines_keeps_the_scope_it_was_handed() {
939        let shared: Arc<dyn Provider> = Arc::new(Declining);
940
941        let scoped = conversation_scoped(Arc::clone(&shared));
942
943        // Not an error path: a provider that does not chain has nothing to
944        // separate, and wrapping it would allocate for nothing.
945        assert!(Arc::ptr_eq(&scoped, &shared));
946    }
947}
mentra/provider.rs

mentra/
provider.rs
