add base model to EmbeddingProvider, not yet leveraged for truncation

KCaverly committed 3 years ago

Commit 2b780ee1 parent d1dec835 files · +57 −2Browse files

Unified · Split

5 files

crates/ai/src/embedding.rs +3 −0
crates/ai/src/providers/dummy.rs +35 −0
crates/ai/src/providers/open_ai/embedding.rs +10 −0
crates/ai/src/providers/open_ai/model.rs +1 −0
crates/semantic_index/src/semantic_index_tests.rs +8 −2

 use rusqlite::ToSql;
 use std::time::Instant;
+use crate::models::LanguageModel;
 #[derive(Debug, PartialEq, Clone)]
 pub struct Embedding(pub Vec<f32>);
 #[async_trait]
 pub trait EmbeddingProvider: Sync + Send {
+    fn base_model(&self) -> Box<dyn LanguageModel>;
     fn is_authenticated(&self) -> bool;
     async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
     fn max_tokens_per_batch(&self) -> usize;

crates/ai/src/providers/dummy.rs+35−0

@@ -3,10 +3,42 @@ use std::time::Instant; · 2 unmodified lines
3	use crate::{	3	use crate::{
4	completion::CompletionRequest,	4	completion::CompletionRequest,
5	embedding::{Embedding, EmbeddingProvider},	5	embedding::{Embedding, EmbeddingProvider},
		6	models::{LanguageModel, TruncationDirection},
6	};	7	};
7	use async_trait::async_trait;	8	use async_trait::async_trait;
8	use serde::Serialize;	9	use serde::Serialize;
9		10
		11	pub struct DummyLanguageModel {}
		12
		13	impl LanguageModel for DummyLanguageModel {
		14	fn name(&self) -> String {
		15	"dummy".to_string()
		16	}
		17	fn capacity(&self) -> anyhow::Result<usize> {
		18	anyhow::Ok(1000)
		19	}
		20	fn truncate(
		21	&self,
		22	content: &str,
		23	length: usize,
		24	direction: crate::models::TruncationDirection,
		25	) -> anyhow::Result<String> {
		26	let truncated = match direction {
		27	TruncationDirection::End => content.chars().collect::<Vec<char>>()[..length]
		28	.iter()
		29	.collect::<String>(),
		30	TruncationDirection::Start => content.chars().collect::<Vec<char>>()[..length]
		31	.iter()
		32	.collect::<String>(),
		33	};
		34
		35	anyhow::Ok(truncated)
		36	}
		37	fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
		38	anyhow::Ok(content.chars().collect::<Vec<char>>().len())
		39	}
		40	}
		41
10	#[derive(Serialize)]	42	#[derive(Serialize)]
11	pub struct DummyCompletionRequest {	43	pub struct DummyCompletionRequest {
12	pub name: String,	44	pub name: String,
@@ -22,6 +54,9 @@ pub struct DummyEmbeddingProvider {} · 9 unmodified lines
22		54
23	#[async_trait]	55	#[async_trait]
24	impl EmbeddingProvider for DummyEmbeddingProvider {	56	impl EmbeddingProvider for DummyEmbeddingProvider {
		57	fn base_model(&self) -> Box<dyn LanguageModel> {
		58	Box::new(DummyLanguageModel {})
		59	}
25	fn is_authenticated(&self) -> bool {	60	fn is_authenticated(&self) -> bool {
26	true	61	true
27	}	62	}

crates/ai/src/providers/open_ai/embedding.rs+10−0

 use util::http::{HttpClient, Request};
 use crate::embedding::{Embedding, EmbeddingProvider};
+use crate::models::LanguageModel;
+use crate::providers::open_ai::OpenAILanguageModel;
 lazy_static! {
     static ref OPENAI_API_KEY: Option<String> = env::var("OPENAI_API_KEY").ok();
 #[derive(Clone)]
 pub struct OpenAIEmbeddingProvider {
+    model: OpenAILanguageModel,
     pub client: Arc<dyn HttpClient>,
     pub executor: Arc<Background>,
     rate_limit_count_rx: watch::Receiver<Option<Instant>>,
         let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
         let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
+        let model = OpenAILanguageModel::load("text-embedding-ada-002");
         OpenAIEmbeddingProvider {
+            model,
             client,
             executor,
             rate_limit_count_rx,
 #[async_trait]
 impl EmbeddingProvider for OpenAIEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(self.model.clone());
+        model
+    }
     fn is_authenticated(&self) -> bool {
         OPENAI_API_KEY.as_ref().is_some()
     }

crates/ai/src/providers/open_ai/model.rs+1−0

@@ -4,6 +4,7 @@ use util::ResultExt; · 3 unmodified lines
4		4
5	use crate::models::{LanguageModel, TruncationDirection};	5	use crate::models::{LanguageModel, TruncationDirection};
6		6
		7	#[derive(Clone)]
7	pub struct OpenAILanguageModel {	8	pub struct OpenAILanguageModel {
8	name: String,	9	name: String,
9	bpe: Option<CoreBPE>,	10	bpe: Option<CoreBPE>,

crates/semantic_index/src/semantic_index_tests.rs+8−2

     semantic_index_settings::SemanticIndexSettings,
     FileToEmbed, JobHandle, SearchResult, SemanticIndex, EMBEDDING_QUEUE_FLUSH_TIMEOUT,
 };
-use ai::embedding::{Embedding, EmbeddingProvider};
+use ai::providers::dummy::{DummyEmbeddingProvider, DummyLanguageModel};
-use ai::providers::dummy::DummyEmbeddingProvider;
+use ai::{
+    embedding::{Embedding, EmbeddingProvider},
+    models::LanguageModel,
+};
 use anyhow::Result;
 use async_trait::async_trait;
 use gpui::{executor::Deterministic, Task, TestAppContext};
 #[async_trait]
 impl EmbeddingProvider for FakeEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        Box::new(DummyLanguageModel {})
+    }
     fn is_authenticated(&self) -> bool {
         true
     }

add base model to EmbeddingProvider, not yet leveraged for truncation

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code