Skip to main content

just_llm_client/
capability.rs

1use std::{
2    fmt,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use async_trait::async_trait;
8use futures_core::Stream;
9
10use crate::{
11    error::{BackendError, Capability, CapabilityError},
12    types::{balance::BalanceSnapshot, chat::ChatCompletionChunk, model::ModelCatalogResponse},
13};
14
15/// Stream of normalized chat-completion chunks.
16///
17/// Wrapper around a `Pin<Box<dyn Stream<...>>>` that implements [`Stream`] so all `StreamExt`
18/// methods (`.next()`, `.map()`, `.collect()`, etc.) work as expected.
19#[must_use = "streams are lazy; call .next() to drive them"]
20pub struct ChatCompletionStream {
21    inner: Pin<
22        Box<
23            dyn Stream<Item = Result<ChatCompletionChunk, just_common::error::TransportError>>
24                + Send,
25        >,
26    >,
27}
28
29impl ChatCompletionStream {
30    /// Wraps a boxed, pinned stream of chunks into a typed [`ChatCompletionStream`].
31    pub fn new(
32        inner: Pin<
33            Box<
34                dyn Stream<Item = Result<ChatCompletionChunk, just_common::error::TransportError>>
35                    + Send,
36            >,
37        >,
38    ) -> Self {
39        Self { inner }
40    }
41}
42
43impl fmt::Debug for ChatCompletionStream {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_struct("ChatCompletionStream")
46            .finish_non_exhaustive()
47    }
48}
49
50impl Stream for ChatCompletionStream {
51    type Item = Result<ChatCompletionChunk, just_common::error::TransportError>;
52
53    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
54        self.inner.as_mut().poll_next(cx)
55    }
56}
57
58/// Root identity trait shared by all client capabilities.
59///
60/// Every backend is identified by its [`family`](Self::family) — the stable string
61/// ("deepseek", "openai-compatible") used for error attribution and diagnostics.
62pub trait Identifiable: Send + Sync {
63    /// Returns the backend family used to identify and attribute this backend in errors.
64    ///
65    /// This is the *instance* family of an already-constructed backend. `LlmBackend` also
66    /// provides a static `family()` (used by `BackendFactory` to key registration before any
67    /// backend exists); both return the same centralized `family` constant.
68    fn family(&self) -> &'static str;
69}
70
71/// List available models from the provider.
72#[async_trait]
73pub trait ModelCatalog: Identifiable {
74    /// Returns the provider's current model catalog.
75    async fn list_models(&self) -> Result<ModelCatalogResponse, BackendError>;
76}
77
78/// Query account balance or quota state.
79#[async_trait]
80pub trait Balance: Identifiable {
81    /// Returns the provider's current balance snapshot.
82    async fn get_balance(&self) -> Result<BalanceSnapshot, BackendError>;
83}
84
85/// Explicit capability negotiation for runtime-selected or otherwise abstract backends.
86///
87/// Each successful negotiation returns a handle that only exposes the requested behavior. If a
88/// backend does not support a capability, [`CapabilityError`] is surfaced here instead of an error
89/// from the capability trait itself.
90pub trait CapabilityNegotiation: Identifiable {
91    /// Returns a handle for model catalog inspection when the backend supports it.
92    fn model_catalog(&self) -> Result<&dyn ModelCatalog, CapabilityError> {
93        Err(CapabilityError::unsupported(
94            self.family(),
95            Capability::ModelCatalog,
96        ))
97    }
98
99    /// Returns a handle for balance inspection when the backend supports it.
100    fn balance(&self) -> Result<&dyn Balance, CapabilityError> {
101        Err(CapabilityError::unsupported(
102            self.family(),
103            Capability::Balance,
104        ))
105    }
106}