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 }}
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
}
}