move api authentication to embedding provider

This commit is contained in:
KCaverly 2023-09-21 14:00:00 -04:00
parent 997f362cc2
commit 7b63369df2
3 changed files with 13 additions and 7 deletions

View File

@ -117,6 +117,7 @@ struct OpenAIEmbeddingUsage {
#[async_trait]
pub trait EmbeddingProvider: Sync + Send {
fn is_authenticated(&self) -> bool;
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
fn max_tokens_per_batch(&self) -> usize;
fn truncate(&self, span: &str) -> (String, usize);
@ -127,6 +128,9 @@ pub struct DummyEmbeddings {}
#[async_trait]
impl EmbeddingProvider for DummyEmbeddings {
fn is_authenticated(&self) -> bool {
true
}
fn rate_limit_expiration(&self) -> Option<Instant> {
None
}
@ -229,6 +233,9 @@ impl OpenAIEmbeddings {
#[async_trait]
impl EmbeddingProvider for OpenAIEmbeddings {
fn is_authenticated(&self) -> bool {
OPENAI_API_KEY.as_ref().is_some()
}
fn max_tokens_per_batch(&self) -> usize {
50000
}

View File

@ -281,12 +281,8 @@ impl SemanticIndex {
settings::get::<SemanticIndexSettings>(cx).enabled
}
pub fn has_api_key(&self) -> bool {
OPENAI_API_KEY.as_ref().is_some()
}
pub fn status(&self, project: &ModelHandle<Project>) -> SemanticIndexStatus {
if !self.has_api_key() {
if !self.embedding_provider.is_authenticated() {
return SemanticIndexStatus::NotAuthenticated;
}
@ -980,8 +976,8 @@ impl SemanticIndex {
project: ModelHandle<Project>,
cx: &mut ModelContext<Self>,
) -> Task<Result<()>> {
if !self.has_api_key() {
return Task::ready(Err(anyhow!("no open ai key present")));
if !self.embedding_provider.is_authenticated() {
return Task::ready(Err(anyhow!("user is not authenticated")));
}
if !self.projects.contains_key(&project.downgrade()) {

View File

@ -1267,6 +1267,9 @@ impl FakeEmbeddingProvider {
#[async_trait]
impl EmbeddingProvider for FakeEmbeddingProvider {
fn is_authenticated(&self) -> bool {
true
}
fn truncate(&self, span: &str) -> (String, usize) {
(span.to_string(), 1)
}