369 lines · 10.5 KB · Rust
RawDownload
use crate::{    AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,    LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,    LanguageModelToolChoice,};use anyhow::anyhow;use futures::{FutureExt, channel::mpsc, future::BoxFuture, stream::BoxStream, stream::StreamExt};use gpui::{App, AsyncApp, Entity, Task};use http_client::Result;use parking_lot::Mutex;use std::sync::{    Arc,    atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},}; #[derive(Clone)]pub struct FakeLanguageModelProvider {    id: LanguageModelProviderId,    name: LanguageModelProviderName,    models: Vec<Arc<dyn LanguageModel>>,} impl Default for FakeLanguageModelProvider {    fn default() -> Self {        Self {            id: LanguageModelProviderId::from("fake".to_string()),            name: LanguageModelProviderName::from("Fake".to_string()),            models: vec![Arc::new(FakeLanguageModel::default())],        }    }} impl LanguageModelProviderState for FakeLanguageModelProvider {    type ObservableEntity = ();     fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {        None    }} impl LanguageModelProvider for FakeLanguageModelProvider {    fn id(&self) -> LanguageModelProviderId {        self.id.clone()    }     fn name(&self) -> LanguageModelProviderName {        self.name.clone()    }     fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {        self.models.first().cloned()    }     fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {        self.models.first().cloned()    }     fn provided_models(&self, _: &App) -> Vec<Arc<dyn LanguageModel>> {        self.models.clone()    }     fn is_authenticated(&self, _: &App) -> bool {        true    }     fn authenticate(&self, _: &mut App) -> Task<Result<(), AuthenticateError>> {        Task::ready(Ok(()))    }     fn settings_view(&self, _: &mut App) -> Option<crate::ProviderSettingsView> {        None    }} impl FakeLanguageModelProvider {    pub fn new(id: LanguageModelProviderId, name: LanguageModelProviderName) -> Self {        Self {            id,            name,            models: vec![Arc::new(FakeLanguageModel::default())],        }    }     pub fn with_models(mut self, models: Vec<Arc<dyn LanguageModel>>) -> Self {        self.models = models;        self    }     pub fn test_model(&self) -> FakeLanguageModel {        FakeLanguageModel::default()    }} #[derive(Debug, PartialEq)]pub struct ToolUseRequest {    pub request: LanguageModelRequest,    pub name: String,    pub description: String,    pub schema: serde_json::Value,} pub struct FakeLanguageModel {    id: LanguageModelId,    name: LanguageModelName,    provider_id: LanguageModelProviderId,    provider_name: LanguageModelProviderName,    current_completion_txs: Mutex<        Vec<(            LanguageModelRequest,            mpsc::UnboundedSender<                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,            >,        )>,    >,    forbid_requests: AtomicBool,    supports_thinking: AtomicBool,    supports_disabling_thinking: AtomicBool,    supports_streaming_tools: AtomicBool,    supports_images: AtomicBool,    supports_server_side_compaction: AtomicBool,    max_token_count: AtomicU64,    max_output_tokens: AtomicU64,} impl Default for FakeLanguageModel {    fn default() -> Self {        Self {            id: LanguageModelId::from("fake".to_string()),            name: LanguageModelName::from("Fake".to_string()),            provider_id: LanguageModelProviderId::from("fake".to_string()),            provider_name: LanguageModelProviderName::from("Fake".to_string()),            current_completion_txs: Mutex::new(Vec::new()),            forbid_requests: AtomicBool::new(false),            supports_thinking: AtomicBool::new(false),            supports_disabling_thinking: AtomicBool::new(true),            supports_streaming_tools: AtomicBool::new(false),            supports_images: AtomicBool::new(false),            supports_server_side_compaction: AtomicBool::new(false),            max_token_count: AtomicU64::new(1_000_000),            max_output_tokens: AtomicU64::new(0),        }    }} impl FakeLanguageModel {    pub fn with_id_and_thinking(        provider_id: &str,        id: &str,        name: &str,        supports_thinking: bool,    ) -> Self {        Self {            id: LanguageModelId::from(id.to_string()),            name: LanguageModelName::from(name.to_string()),            provider_id: LanguageModelProviderId::from(provider_id.to_string()),            supports_thinking: AtomicBool::new(supports_thinking),            ..Default::default()        }    }     pub fn allow_requests(&self) {        self.forbid_requests.store(false, SeqCst);    }     pub fn forbid_requests(&self) {        self.forbid_requests.store(true, SeqCst);    }     pub fn set_supports_thinking(&self, supports: bool) {        self.supports_thinking.store(supports, SeqCst);    }     pub fn set_supports_disabling_thinking(&self, supports: bool) {        self.supports_disabling_thinking.store(supports, SeqCst);    }     pub fn set_supports_streaming_tools(&self, supports: bool) {        self.supports_streaming_tools.store(supports, SeqCst);    }     pub fn set_supports_images(&self, supports: bool) {        self.supports_images.store(supports, SeqCst);    }     pub fn set_supports_server_side_compaction(&self, supports: bool) {        self.supports_server_side_compaction.store(supports, SeqCst);    }     pub fn set_max_token_count(&self, count: u64) {        self.max_token_count.store(count, SeqCst);    }     pub fn set_max_output_tokens(&self, count: Option<u64>) {        self.max_output_tokens            .store(count.unwrap_or_default(), SeqCst);    }     pub fn pending_completions(&self) -> Vec<LanguageModelRequest> {        self.current_completion_txs            .lock()            .iter()            .map(|(request, _)| request.clone())            .collect()    }     pub fn completion_count(&self) -> usize {        self.current_completion_txs.lock().len()    }     pub fn send_completion_stream_text_chunk(        &self,        request: &LanguageModelRequest,        chunk: impl Into<String>,    ) {        self.send_completion_stream_event(            request,            LanguageModelCompletionEvent::Text(chunk.into()),        );    }     pub fn send_completion_stream_event(        &self,        request: &LanguageModelRequest,        event: impl Into<LanguageModelCompletionEvent>,    ) {        let current_completion_txs = self.current_completion_txs.lock();        let tx = current_completion_txs            .iter()            .find(|(req, _)| req == request)            .map(|(_, tx)| tx)            .unwrap();        tx.unbounded_send(Ok(event.into())).unwrap();    }     pub fn send_completion_stream_error(        &self,        request: &LanguageModelRequest,        error: impl Into<LanguageModelCompletionError>,    ) {        let current_completion_txs = self.current_completion_txs.lock();        let tx = current_completion_txs            .iter()            .find(|(req, _)| req == request)            .map(|(_, tx)| tx)            .unwrap();        tx.unbounded_send(Err(error.into())).unwrap();    }     pub fn end_completion_stream(&self, request: &LanguageModelRequest) {        self.current_completion_txs            .lock()            .retain(|(req, _)| req != request);    }     pub fn send_last_completion_stream_text_chunk(&self, chunk: impl Into<String>) {        self.send_completion_stream_text_chunk(self.pending_completions().last().unwrap(), chunk);    }     pub fn send_last_completion_stream_event(        &self,        event: impl Into<LanguageModelCompletionEvent>,    ) {        self.send_completion_stream_event(self.pending_completions().last().unwrap(), event);    }     pub fn send_last_completion_stream_error(        &self,        error: impl Into<LanguageModelCompletionError>,    ) {        self.send_completion_stream_error(self.pending_completions().last().unwrap(), error);    }     pub fn end_last_completion_stream(&self) {        self.end_completion_stream(self.pending_completions().last().unwrap());    }} impl LanguageModel for FakeLanguageModel {    fn id(&self) -> LanguageModelId {        self.id.clone()    }     fn name(&self) -> LanguageModelName {        self.name.clone()    }     fn provider_id(&self) -> LanguageModelProviderId {        self.provider_id.clone()    }     fn provider_name(&self) -> LanguageModelProviderName {        self.provider_name.clone()    }     fn supports_tools(&self) -> bool {        false    }     fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {        false    }     fn supports_images(&self) -> bool {        self.supports_images.load(SeqCst)    }     fn supports_server_side_compaction(&self) -> bool {        self.supports_server_side_compaction.load(SeqCst)    }     fn supports_thinking(&self) -> bool {        self.supports_thinking.load(SeqCst)    }     fn supports_disabling_thinking(&self) -> bool {        self.supports_disabling_thinking.load(SeqCst)    }     fn supports_streaming_tools(&self) -> bool {        self.supports_streaming_tools.load(SeqCst)    }     fn telemetry_id(&self) -> String {        "fake".to_string()    }     fn max_token_count(&self) -> u64 {        self.max_token_count.load(SeqCst)    }     fn max_output_tokens(&self) -> Option<u64> {        let max_output_tokens = self.max_output_tokens.load(SeqCst);        if max_output_tokens == 0 {            None        } else {            Some(max_output_tokens)        }    }     fn stream_completion(        &self,        request: LanguageModelRequest,        _: &AsyncApp,    ) -> BoxFuture<        'static,        Result<            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,            LanguageModelCompletionError,        >,    > {        if self.forbid_requests.load(SeqCst) {            async move {                Err(LanguageModelCompletionError::Other(anyhow!(                    "requests are forbidden"                )))            }            .boxed()        } else {            let (tx, rx) = mpsc::unbounded();            self.current_completion_txs.lock().push((request, tx));            async move { Ok(rx.boxed()) }.boxed()        }    }     fn as_fake(&self) -> &Self {        self    }}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t