Unified · Split
| @@ -5,6 +5,8 @@ use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef}; · 4 unmodified lines | ||
| 5 | 5 | use rusqlite::ToSql; |
| 6 | 6 | use std::time::Instant; |
| 7 | 7 | |
| 8 | +use crate::models::LanguageModel; | |
| 9 | + | |
| 8 | 10 | #[derive(Debug, PartialEq, Clone)] |
| 9 | 11 | pub struct Embedding(pub Vec<f32>); |
| 10 | 12 | |
| @@ -66,6 +68,7 @@ impl Embedding { · 55 unmodified lines | ||
| 66 | 68 | |
| 67 | 69 | #[async_trait] |
| 68 | 70 | pub trait EmbeddingProvider: Sync + Send { |
| 71 | + fn base_model(&self) -> Box<dyn LanguageModel>; | |
| 69 | 72 | fn is_authenticated(&self) -> bool; |
| 70 | 73 | async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>; |
| 71 | 74 | fn max_tokens_per_batch(&self) -> usize; |
| @@ -3,10 +3,42 @@ use std::time::Instant; · 2 unmodified lines | ||
| 3 | 3 | use crate::{ |
| 4 | 4 | completion::CompletionRequest, |
| 5 | 5 | embedding::{Embedding, EmbeddingProvider}, |
| 6 | + models::{LanguageModel, TruncationDirection}, | |
| 6 | 7 | }; |
| 7 | 8 | use async_trait::async_trait; |
| 8 | 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 | 42 | #[derive(Serialize)] |
| 11 | 43 | pub struct DummyCompletionRequest { |
| 12 | 44 | pub name: String, |
| @@ -22,6 +54,9 @@ pub struct DummyEmbeddingProvider {} · 9 unmodified lines | ||
| 22 | 54 | |
| 23 | 55 | #[async_trait] |
| 24 | 56 | impl EmbeddingProvider for DummyEmbeddingProvider { |
| 57 | + fn base_model(&self) -> Box<dyn LanguageModel> { | |
| 58 | + Box::new(DummyLanguageModel {}) | |
| 59 | + } | |
| 25 | 60 | fn is_authenticated(&self) -> bool { |
| 26 | 61 | true |
| 27 | 62 | } |
| @@ -19,6 +19,8 @@ use tiktoken_rs::{cl100k_base, CoreBPE}; · 18 unmodified lines | ||
| 19 | 19 | use util::http::{HttpClient, Request}; |
| 20 | 20 | |
| 21 | 21 | use crate::embedding::{Embedding, EmbeddingProvider}; |
| 22 | +use crate::models::LanguageModel; | |
| 23 | +use crate::providers::open_ai::OpenAILanguageModel; | |
| 22 | 24 | |
| 23 | 25 | lazy_static! { |
| 24 | 26 | static ref OPENAI_API_KEY: Option<String> = env::var("OPENAI_API_KEY").ok(); |
| @@ -27,6 +29,7 @@ lazy_static! { · 2 unmodified lines | ||
| 27 | 29 | |
| 28 | 30 | #[derive(Clone)] |
| 29 | 31 | pub struct OpenAIEmbeddingProvider { |
| 32 | + model: OpenAILanguageModel, | |
| 30 | 33 | pub client: Arc<dyn HttpClient>, |
| 31 | 34 | pub executor: Arc<Background>, |
| 32 | 35 | rate_limit_count_rx: watch::Receiver<Option<Instant>>, |
| @@ -65,7 +68,10 @@ impl OpenAIEmbeddingProvider { · 32 unmodified lines | ||
| 65 | 68 | let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None); |
| 66 | 69 | let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx)); |
| 67 | 70 | |
| 71 | + let model = OpenAILanguageModel::load("text-embedding-ada-002"); | |
| 72 | + | |
| 68 | 73 | OpenAIEmbeddingProvider { |
| 74 | + model, | |
| 69 | 75 | client, |
| 70 | 76 | executor, |
| 71 | 77 | rate_limit_count_rx, |
| @@ -131,6 +137,10 @@ impl OpenAIEmbeddingProvider { · 59 unmodified lines | ||
| 131 | 137 | |
| 132 | 138 | #[async_trait] |
| 133 | 139 | impl EmbeddingProvider for OpenAIEmbeddingProvider { |
| 140 | + fn base_model(&self) -> Box<dyn LanguageModel> { | |
| 141 | + let model: Box<dyn LanguageModel> = Box::new(self.model.clone()); | |
| 142 | + model | |
| 143 | + } | |
| 134 | 144 | fn is_authenticated(&self) -> bool { |
| 135 | 145 | OPENAI_API_KEY.as_ref().is_some() |
| 136 | 146 | } |
| @@ -4,6 +4,7 @@ use util::ResultExt; · 3 unmodified lines | ||
| 4 | 4 | |
| 5 | 5 | use crate::models::{LanguageModel, TruncationDirection}; |
| 6 | 6 | |
| 7 | +#[derive(Clone)] | |
| 7 | 8 | pub struct OpenAILanguageModel { |
| 8 | 9 | name: String, |
| 9 | 10 | bpe: Option<CoreBPE>, |
| @@ -4,8 +4,11 @@ use crate::{ · 3 unmodified lines | ||
| 4 | 4 | semantic_index_settings::SemanticIndexSettings, |
| 5 | 5 | FileToEmbed, JobHandle, SearchResult, SemanticIndex, EMBEDDING_QUEUE_FLUSH_TIMEOUT, |
| 6 | 6 | }; |
| 7 | −use ai::embedding::{Embedding, EmbeddingProvider}; | |
| 8 | −use ai::providers::dummy::DummyEmbeddingProvider; | |
| 7 | +use ai::providers::dummy::{DummyEmbeddingProvider, DummyLanguageModel}; | |
| 8 | +use ai::{ | |
| 9 | + embedding::{Embedding, EmbeddingProvider}, | |
| 10 | + models::LanguageModel, | |
| 11 | +}; | |
| 9 | 12 | use anyhow::Result; |
| 10 | 13 | use async_trait::async_trait; |
| 11 | 14 | use gpui::{executor::Deterministic, Task, TestAppContext}; |
| @@ -1282,6 +1285,9 @@ impl FakeEmbeddingProvider { · 1270 unmodified lines | ||
| 1282 | 1285 | |
| 1283 | 1286 | #[async_trait] |
| 1284 | 1287 | impl EmbeddingProvider for FakeEmbeddingProvider { |
| 1288 | + fn base_model(&self) -> Box<dyn LanguageModel> { | |
| 1289 | + Box::new(DummyLanguageModel {}) | |
| 1290 | + } | |
| 1285 | 1291 | fn is_authenticated(&self) -> bool { |
| 1286 | 1292 | true |
| 1287 | 1293 | } |