1use aisdk::core::capabilities::{
4 ReasoningSupport, StructuredOutputSupport, TextInputSupport, TextOutputSupport,
5 ToolCallSupport,
6};
7use aisdk::core::language_model::{
8 LanguageModel, LanguageModelOptions, LanguageModelResponse, LanguageModelStreamChunk,
9};
10use aisdk::Result as AisdkResult;
11use async_trait::async_trait;
12use futures_util::Stream;
13use std::fmt;
14use std::pin::Pin;
15use std::sync::Arc;
16use tokio::sync::Mutex;
17
18type ProviderStream =
19 Pin<Box<dyn Stream<Item = AisdkResult<Vec<LanguageModelStreamChunk>>> + Send>>;
20
21#[async_trait]
22trait DynModel: Send + Sync + fmt::Debug {
23 fn name(&self) -> String;
24 async fn generate_text(&self, options: LanguageModelOptions)
25 -> AisdkResult<LanguageModelResponse>;
26 async fn stream_text(&self, options: LanguageModelOptions) -> AisdkResult<ProviderStream>;
27}
28
29struct Holding<M: LanguageModel> {
30 name: String,
31 inner: Mutex<M>,
32}
33
34impl<M: LanguageModel> fmt::Debug for Holding<M> {
35 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36 f.debug_struct("Holding").field("name", &self.name).finish()
37 }
38}
39
40#[async_trait]
41impl<M: LanguageModel> DynModel for Holding<M> {
42 fn name(&self) -> String {
43 self.name.clone()
44 }
45
46 async fn generate_text(
47 &self,
48 options: LanguageModelOptions,
49 ) -> AisdkResult<LanguageModelResponse> {
50 self.inner.lock().await.generate_text(options).await
51 }
52
53 async fn stream_text(&self, options: LanguageModelOptions) -> AisdkResult<ProviderStream> {
54 self.inner.lock().await.stream_text(options).await
55 }
56}
57
58#[derive(Clone)]
63pub struct SharedModel {
64 inner: Arc<dyn DynModel>,
65}
66
67impl fmt::Debug for SharedModel {
68 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69 f.debug_struct("SharedModel")
70 .field("name", &self.inner.name())
71 .finish()
72 }
73}
74
75impl SharedModel {
76 pub fn wrap<M: LanguageModel>(model: M) -> Self {
78 let name = model.name();
79 Self {
80 inner: Arc::new(Holding {
81 name,
82 inner: Mutex::new(model),
83 }),
84 }
85 }
86}
87
88impl TextInputSupport for SharedModel {}
89impl TextOutputSupport for SharedModel {}
90impl ToolCallSupport for SharedModel {}
91impl StructuredOutputSupport for SharedModel {}
92impl ReasoningSupport for SharedModel {}
93
94#[async_trait]
95impl LanguageModel for SharedModel {
96 fn name(&self) -> String {
97 self.inner.name()
98 }
99
100 async fn generate_text(
101 &mut self,
102 options: LanguageModelOptions,
103 ) -> AisdkResult<LanguageModelResponse> {
104 self.inner.generate_text(options).await
105 }
106
107 async fn stream_text(&mut self, options: LanguageModelOptions) -> AisdkResult<ProviderStream> {
108 self.inner.stream_text(options).await
109 }
110}