just_llm_client/
capability.rs1use 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#[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 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
58pub trait Identifiable: Send + Sync {
63 fn family(&self) -> &'static str;
69}
70
71#[async_trait]
73pub trait ModelCatalog: Identifiable {
74 async fn list_models(&self) -> Result<ModelCatalogResponse, BackendError>;
76}
77
78#[async_trait]
80pub trait Balance: Identifiable {
81 async fn get_balance(&self) -> Result<BalanceSnapshot, BackendError>;
83}
84
85pub trait CapabilityNegotiation: Identifiable {
91 fn model_catalog(&self) -> Result<&dyn ModelCatalog, CapabilityError> {
93 Err(CapabilityError::unsupported(
94 self.family(),
95 Capability::ModelCatalog,
96 ))
97 }
98
99 fn balance(&self) -> Result<&dyn Balance, CapabilityError> {
101 Err(CapabilityError::unsupported(
102 self.family(),
103 Capability::Balance,
104 ))
105 }
106}