also allow api_key

This commit is contained in:
Louis Dureuil 2024-03-25 11:50:00 +01:00
parent 4136630ea5
commit 817ccc089a
No known key found for this signature in database
4 changed files with 14 additions and 13 deletions

View File

@ -1278,7 +1278,6 @@ pub fn validate_embedding_settings(
// Dimensions get inferred, only model name is required // Dimensions get inferred, only model name is required
check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?;
check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?; check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?;
check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;
check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?;

View File

@ -201,8 +201,8 @@ impl EmbedderOptions {
Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key))
} }
pub fn ollama(url: Option<String>) -> Self { pub fn ollama(api_key: Option<String>, url: Option<String>) -> Self {
Self::Ollama(ollama::EmbedderOptions::with_default_model(url)) Self::Ollama(ollama::EmbedderOptions::with_default_model(api_key, url))
} }
} }

View File

@ -13,11 +13,12 @@ pub struct Embedder {
pub struct EmbedderOptions { pub struct EmbedderOptions {
pub embedding_model: String, pub embedding_model: String,
pub url: Option<String>, pub url: Option<String>,
pub api_key: Option<String>,
} }
impl EmbedderOptions { impl EmbedderOptions {
pub fn with_default_model(url: Option<String>) -> Self { pub fn with_default_model(api_key: Option<String>, url: Option<String>) -> Self {
Self { embedding_model: "nomic-embed-text".into(), url } Self { embedding_model: "nomic-embed-text".into(), api_key, url }
} }
} }
@ -25,7 +26,7 @@ impl Embedder {
pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> { pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
let model = options.embedding_model.as_str(); let model = options.embedding_model.as_str();
let rest_embedder = match RestEmbedder::new(RestEmbedderOptions { let rest_embedder = match RestEmbedder::new(RestEmbedderOptions {
api_key: None, api_key: options.api_key,
distribution: None, distribution: None,
dimensions: None, dimensions: None,
url: options.url.unwrap_or_else(get_ollama_path), url: options.url.unwrap_or_else(get_ollama_path),

View File

@ -114,7 +114,9 @@ impl EmbeddingSettings {
&[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama]
} }
Self::REVISION => &[EmbedderSource::HuggingFace], Self::REVISION => &[EmbedderSource::HuggingFace],
Self::API_KEY => &[EmbedderSource::OpenAi, EmbedderSource::Rest], Self::API_KEY => {
&[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest]
}
Self::DIMENSIONS => { Self::DIMENSIONS => {
&[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest] &[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest]
} }
@ -147,7 +149,7 @@ impl EmbeddingSettings {
&[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE]
} }
EmbedderSource::Ollama => { EmbedderSource::Ollama => {
&[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL] &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL, Self::API_KEY]
} }
EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS],
EmbedderSource::Rest => &[ EmbedderSource::Rest => &[
@ -389,15 +391,14 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
} }
EmbedderSource::Ollama => { EmbedderSource::Ollama => {
let mut options: ollama::EmbedderOptions = let mut options: ollama::EmbedderOptions =
super::ollama::EmbedderOptions::with_default_model(None); super::ollama::EmbedderOptions::with_default_model(
api_key.set(),
url.set(),
);
if let Some(model) = model.set() { if let Some(model) = model.set() {
options.embedding_model = model; options.embedding_model = model;
} }
if let Some(url) = url.set() {
options.url = Some(url)
}
this.embedder_options = super::EmbedderOptions::Ollama(options); this.embedder_options = super::EmbedderOptions::Ollama(options);
} }
EmbedderSource::HuggingFace => { EmbedderSource::HuggingFace => {