language_model.rs

  1mod model;
  2pub mod provider;
  3mod rate_limiter;
  4mod registry;
  5mod request;
  6mod role;
  7pub mod settings;
  8
  9use anyhow::Result;
 10use client::{Client, UserStore};
 11use futures::{future::BoxFuture, stream::BoxStream};
 12use gpui::{
 13    AnyView, AppContext, AsyncAppContext, FocusHandle, Model, SharedString, Task, WindowContext,
 14};
 15pub use model::*;
 16use project::Fs;
 17pub(crate) use rate_limiter::*;
 18pub use registry::*;
 19pub use request::*;
 20pub use role::*;
 21use schemars::JsonSchema;
 22use serde::de::DeserializeOwned;
 23use std::{future::Future, sync::Arc};
 24
 25pub fn init(
 26    user_store: Model<UserStore>,
 27    client: Arc<Client>,
 28    fs: Arc<dyn Fs>,
 29    cx: &mut AppContext,
 30) {
 31    settings::init(fs, cx);
 32    registry::init(user_store, client, cx);
 33}
 34
 35pub trait LanguageModel: Send + Sync {
 36    fn id(&self) -> LanguageModelId;
 37    fn name(&self) -> LanguageModelName;
 38    fn provider_id(&self) -> LanguageModelProviderId;
 39    fn provider_name(&self) -> LanguageModelProviderName;
 40    fn telemetry_id(&self) -> String;
 41
 42    fn max_token_count(&self) -> usize;
 43
 44    fn count_tokens(
 45        &self,
 46        request: LanguageModelRequest,
 47        cx: &AppContext,
 48    ) -> BoxFuture<'static, Result<usize>>;
 49
 50    fn stream_completion(
 51        &self,
 52        request: LanguageModelRequest,
 53        cx: &AsyncAppContext,
 54    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
 55
 56    fn use_any_tool(
 57        &self,
 58        request: LanguageModelRequest,
 59        name: String,
 60        description: String,
 61        schema: serde_json::Value,
 62        cx: &AsyncAppContext,
 63    ) -> BoxFuture<'static, Result<serde_json::Value>>;
 64}
 65
 66impl dyn LanguageModel {
 67    pub fn use_tool<T: LanguageModelTool>(
 68        &self,
 69        request: LanguageModelRequest,
 70        cx: &AsyncAppContext,
 71    ) -> impl 'static + Future<Output = Result<T>> {
 72        let schema = schemars::schema_for!(T);
 73        let schema_json = serde_json::to_value(&schema).unwrap();
 74        let request = self.use_any_tool(request, T::name(), T::description(), schema_json, cx);
 75        async move {
 76            let response = request.await?;
 77            Ok(serde_json::from_value(response)?)
 78        }
 79    }
 80}
 81
 82pub trait LanguageModelTool: 'static + DeserializeOwned + JsonSchema {
 83    fn name() -> String;
 84    fn description() -> String;
 85}
 86
 87pub trait LanguageModelProvider: 'static {
 88    fn id(&self) -> LanguageModelProviderId;
 89    fn name(&self) -> LanguageModelProviderName;
 90    fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>>;
 91    fn load_model(&self, _model: Arc<dyn LanguageModel>, _cx: &AppContext) {}
 92    fn is_authenticated(&self, cx: &AppContext) -> bool;
 93    fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>>;
 94    fn configuration_view(&self, cx: &mut WindowContext) -> (AnyView, Option<FocusHandle>);
 95    fn reset_credentials(&self, cx: &mut AppContext) -> Task<Result<()>>;
 96}
 97
 98pub trait LanguageModelProviderState: 'static {
 99    type ObservableEntity;
100
101    fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>>;
102
103    fn subscribe<T: 'static>(
104        &self,
105        cx: &mut gpui::ModelContext<T>,
106        callback: impl Fn(&mut T, &mut gpui::ModelContext<T>) + 'static,
107    ) -> Option<gpui::Subscription> {
108        let entity = self.observable_entity()?;
109        Some(cx.observe(&entity, move |this, _, cx| {
110            callback(this, cx);
111        }))
112    }
113}
114
115#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
116pub struct LanguageModelId(pub SharedString);
117
118#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
119pub struct LanguageModelName(pub SharedString);
120
121#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
122pub struct LanguageModelProviderId(pub SharedString);
123
124#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
125pub struct LanguageModelProviderName(pub SharedString);
126
127impl From<String> for LanguageModelId {
128    fn from(value: String) -> Self {
129        Self(SharedString::from(value))
130    }
131}
132
133impl From<String> for LanguageModelName {
134    fn from(value: String) -> Self {
135        Self(SharedString::from(value))
136    }
137}
138
139impl From<String> for LanguageModelProviderId {
140    fn from(value: String) -> Self {
141        Self(SharedString::from(value))
142    }
143}
144
145impl From<String> for LanguageModelProviderName {
146    fn from(value: String) -> Self {
147        Self(SharedString::from(value))
148    }
149}