tenant.openagents/omega

1use crate::{
2    AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
3    LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
4    LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
5    LanguageModelToolChoice,
6};
7use anyhow::anyhow;
8use futures::{FutureExt, channel::mpsc, future::BoxFuture, stream::BoxStream, stream::StreamExt};
9use gpui::{App, AsyncApp, Entity, Task};
10use http_client::Result;
11use parking_lot::Mutex;
12use std::sync::{
13    Arc,
14    atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},
15};
16
17#[derive(Clone)]
18pub struct FakeLanguageModelProvider {
19    id: LanguageModelProviderId,
20    name: LanguageModelProviderName,
21    models: Vec<Arc<dyn LanguageModel>>,
22}
23
24impl Default for FakeLanguageModelProvider {
25    fn default() -> Self {
26        Self {
27            id: LanguageModelProviderId::from("fake".to_string()),
28            name: LanguageModelProviderName::from("Fake".to_string()),
29            models: vec![Arc::new(FakeLanguageModel::default())],
30        }
31    }
32}
33
34impl LanguageModelProviderState for FakeLanguageModelProvider {
35    type ObservableEntity = ();
36
37    fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
38        None
39    }
40}
41
42impl LanguageModelProvider for FakeLanguageModelProvider {
43    fn id(&self) -> LanguageModelProviderId {
44        self.id.clone()
45    }
46
47    fn name(&self) -> LanguageModelProviderName {
48        self.name.clone()
49    }
50
51    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
52        self.models.first().cloned()
53    }
54
55    fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
56        self.models.first().cloned()
57    }
58
59    fn provided_models(&self, _: &App) -> Vec<Arc<dyn LanguageModel>> {
60        self.models.clone()
61    }
62
63    fn is_authenticated(&self, _: &App) -> bool {
64        true
65    }
66
67    fn authenticate(&self, _: &mut App) -> Task<Result<(), AuthenticateError>> {
68        Task::ready(Ok(()))
69    }
70
71    fn settings_view(&self, _: &mut App) -> Option<crate::ProviderSettingsView> {
72        None
73    }
74}
75
76impl FakeLanguageModelProvider {
77    pub fn new(id: LanguageModelProviderId, name: LanguageModelProviderName) -> Self {
78        Self {
79            id,
80            name,
81            models: vec![Arc::new(FakeLanguageModel::default())],
82        }
83    }
84
85    pub fn with_models(mut self, models: Vec<Arc<dyn LanguageModel>>) -> Self {
86        self.models = models;
87        self
88    }
89
90    pub fn test_model(&self) -> FakeLanguageModel {
91        FakeLanguageModel::default()
92    }
93}
94
95#[derive(Debug, PartialEq)]
96pub struct ToolUseRequest {
97    pub request: LanguageModelRequest,
98    pub name: String,
99    pub description: String,
100    pub schema: serde_json::Value,
101}
102
103pub struct FakeLanguageModel {
104    id: LanguageModelId,
105    name: LanguageModelName,
106    provider_id: LanguageModelProviderId,
107    provider_name: LanguageModelProviderName,
108    current_completion_txs: Mutex<
109        Vec<(
110            LanguageModelRequest,
111            mpsc::UnboundedSender<
112                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
113            >,
114        )>,
115    >,
116    forbid_requests: AtomicBool,
117    supports_thinking: AtomicBool,
118    supports_disabling_thinking: AtomicBool,
119    supports_streaming_tools: AtomicBool,
120    supports_images: AtomicBool,
121    supports_server_side_compaction: AtomicBool,
122    max_token_count: AtomicU64,
123    max_output_tokens: AtomicU64,
124}
125
126impl Default for FakeLanguageModel {
127    fn default() -> Self {
128        Self {
129            id: LanguageModelId::from("fake".to_string()),
130            name: LanguageModelName::from("Fake".to_string()),
131            provider_id: LanguageModelProviderId::from("fake".to_string()),
132            provider_name: LanguageModelProviderName::from("Fake".to_string()),
133            current_completion_txs: Mutex::new(Vec::new()),
134            forbid_requests: AtomicBool::new(false),
135            supports_thinking: AtomicBool::new(false),
136            supports_disabling_thinking: AtomicBool::new(true),
137            supports_streaming_tools: AtomicBool::new(false),
138            supports_images: AtomicBool::new(false),
139            supports_server_side_compaction: AtomicBool::new(false),
140            max_token_count: AtomicU64::new(1_000_000),
141            max_output_tokens: AtomicU64::new(0),
142        }
143    }
144}
145
146impl FakeLanguageModel {
147    pub fn with_id_and_thinking(
148        provider_id: &str,
149        id: &str,
150        name: &str,
151        supports_thinking: bool,
152    ) -> Self {
153        Self {
154            id: LanguageModelId::from(id.to_string()),
155            name: LanguageModelName::from(name.to_string()),
156            provider_id: LanguageModelProviderId::from(provider_id.to_string()),
157            supports_thinking: AtomicBool::new(supports_thinking),
158            ..Default::default()
159        }
160    }
161
162    pub fn allow_requests(&self) {
163        self.forbid_requests.store(false, SeqCst);
164    }
165
166    pub fn forbid_requests(&self) {
167        self.forbid_requests.store(true, SeqCst);
168    }
169
170    pub fn set_supports_thinking(&self, supports: bool) {
171        self.supports_thinking.store(supports, SeqCst);
172    }
173
174    pub fn set_supports_disabling_thinking(&self, supports: bool) {
175        self.supports_disabling_thinking.store(supports, SeqCst);
176    }
177
178    pub fn set_supports_streaming_tools(&self, supports: bool) {
179        self.supports_streaming_tools.store(supports, SeqCst);
180    }
181
182    pub fn set_supports_images(&self, supports: bool) {
183        self.supports_images.store(supports, SeqCst);
184    }
185
186    pub fn set_supports_server_side_compaction(&self, supports: bool) {
187        self.supports_server_side_compaction.store(supports, SeqCst);
188    }
189
190    pub fn set_max_token_count(&self, count: u64) {
191        self.max_token_count.store(count, SeqCst);
192    }
193
194    pub fn set_max_output_tokens(&self, count: Option<u64>) {
195        self.max_output_tokens
196            .store(count.unwrap_or_default(), SeqCst);
197    }
198
199    pub fn pending_completions(&self) -> Vec<LanguageModelRequest> {
200        self.current_completion_txs
201            .lock()
202            .iter()
203            .map(|(request, _)| request.clone())
204            .collect()
205    }
206
207    pub fn completion_count(&self) -> usize {
208        self.current_completion_txs.lock().len()
209    }
210
211    pub fn send_completion_stream_text_chunk(
212        &self,
213        request: &LanguageModelRequest,
214        chunk: impl Into<String>,
215    ) {
216        self.send_completion_stream_event(
217            request,
218            LanguageModelCompletionEvent::Text(chunk.into()),
219        );
220    }
221
222    pub fn send_completion_stream_event(
223        &self,
224        request: &LanguageModelRequest,
225        event: impl Into<LanguageModelCompletionEvent>,
226    ) {
227        let current_completion_txs = self.current_completion_txs.lock();
228        let tx = current_completion_txs
229            .iter()
230            .find(|(req, _)| req == request)
231            .map(|(_, tx)| tx)
232            .unwrap();
233        tx.unbounded_send(Ok(event.into())).unwrap();
234    }
235
236    pub fn send_completion_stream_error(
237        &self,
238        request: &LanguageModelRequest,
239        error: impl Into<LanguageModelCompletionError>,
240    ) {
241        let current_completion_txs = self.current_completion_txs.lock();
242        let tx = current_completion_txs
243            .iter()
244            .find(|(req, _)| req == request)
245            .map(|(_, tx)| tx)
246            .unwrap();
247        tx.unbounded_send(Err(error.into())).unwrap();
248    }
249
250    pub fn end_completion_stream(&self, request: &LanguageModelRequest) {
251        self.current_completion_txs
252            .lock()
253            .retain(|(req, _)| req != request);
254    }
255
256    pub fn send_last_completion_stream_text_chunk(&self, chunk: impl Into<String>) {
257        self.send_completion_stream_text_chunk(self.pending_completions().last().unwrap(), chunk);
258    }
259
260    pub fn send_last_completion_stream_event(
261        &self,
262        event: impl Into<LanguageModelCompletionEvent>,
263    ) {
264        self.send_completion_stream_event(self.pending_completions().last().unwrap(), event);
265    }
266
267    pub fn send_last_completion_stream_error(
268        &self,
269        error: impl Into<LanguageModelCompletionError>,
270    ) {
271        self.send_completion_stream_error(self.pending_completions().last().unwrap(), error);
272    }
273
274    pub fn end_last_completion_stream(&self) {
275        self.end_completion_stream(self.pending_completions().last().unwrap());
276    }
277}
278
279impl LanguageModel for FakeLanguageModel {
280    fn id(&self) -> LanguageModelId {
281        self.id.clone()
282    }
283
284    fn name(&self) -> LanguageModelName {
285        self.name.clone()
286    }
287
288    fn provider_id(&self) -> LanguageModelProviderId {
289        self.provider_id.clone()
290    }
291
292    fn provider_name(&self) -> LanguageModelProviderName {
293        self.provider_name.clone()
294    }
295
296    fn supports_tools(&self) -> bool {
297        false
298    }
299
300    fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
301        false
302    }
303
304    fn supports_images(&self) -> bool {
305        self.supports_images.load(SeqCst)
306    }
307
308    fn supports_server_side_compaction(&self) -> bool {
309        self.supports_server_side_compaction.load(SeqCst)
310    }
311
312    fn supports_thinking(&self) -> bool {
313        self.supports_thinking.load(SeqCst)
314    }
315
316    fn supports_disabling_thinking(&self) -> bool {
317        self.supports_disabling_thinking.load(SeqCst)
318    }
319
320    fn supports_streaming_tools(&self) -> bool {
321        self.supports_streaming_tools.load(SeqCst)
322    }
323
324    fn telemetry_id(&self) -> String {
325        "fake".to_string()
326    }
327
328    fn max_token_count(&self) -> u64 {
329        self.max_token_count.load(SeqCst)
330    }
331
332    fn max_output_tokens(&self) -> Option<u64> {
333        let max_output_tokens = self.max_output_tokens.load(SeqCst);
334        if max_output_tokens == 0 {
335            None
336        } else {
337            Some(max_output_tokens)
338        }
339    }
340
341    fn stream_completion(
342        &self,
343        request: LanguageModelRequest,
344        _: &AsyncApp,
345    ) -> BoxFuture<
346        'static,
347        Result<
348            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
349            LanguageModelCompletionError,
350        >,
351    > {
352        if self.forbid_requests.load(SeqCst) {
353            async move {
354                Err(LanguageModelCompletionError::Other(anyhow!(
355                    "requests are forbidden"
356                )))
357            }
358            .boxed()
359        } else {
360            let (tx, rx) = mpsc::unbounded();
361            self.current_completion_txs.lock().push((request, tx));
362            async move { Ok(rx.boxed()) }.boxed()
363        }
364    }
365
366    fn as_fake(&self) -> &Self {
367        self
368    }
369}
370