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
check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, 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(&query, EmbeddingSettings::QUERY, inferred_source, name)?;

View File

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

View File

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

View File

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